Fix shape mismatch when train only with CPU

swook · Sep 28, 2018 · e6bcefb · e6bcefb
1 parent cfb31f6
commit e6bcefb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/src/datasources/unityeyes.py b/src/datasources/unityeyes.py
@@ -296,6 +296,6 @@ def noisy_value_from_type(augmentation_type):
                 for landmark in entry['landmarks']
             ]).astype(np.float32)
             if self.data_format == 'NHWC':
-                np.transpose(entry['heatmaps'], (1, 2, 0))
+                entry['heatmaps'] = np.transpose(entry['heatmaps'], (1, 2, 0))
 
         return entry
diff --git a/src/models/elg.py b/src/models/elg.py
@@ -75,7 +75,10 @@ def build_model(self, data_sources: Dict[str, BaseDataSource], mode: str):
         with tf.variable_scope('hourglass'):
             # TODO: Find better way to specify no. landmarks
             if y1 is not None:
-                self._hg_num_landmarks = y1.shape.as_list()[1]
+                if self._data_format == 'NCHW':
+		    self._hg_num_landmarks = y1.shape.as_list()[1]
+                if self._data_format == 'NHWC':
+		    self._hg_num_landmarks = y1.shape.as_list()[3]
             else:
                 self._hg_num_landmarks = 18
             assert self._hg_num_landmarks == 18