fix acc metric shape miss match

apache · Apr 6, 2018 · 4232318 · 4232318
1 parent 590ec3f
commit 4232318
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/python/mxnet/metric.py b/python/mxnet/metric.py
@@ -417,11 +417,14 @@ def update(self, labels, preds):
                 pred_label = ndarray.argmax(pred_label, axis=self.axis)
             pred_label = pred_label.asnumpy().astype('int32')
             label = label.asnumpy().astype('int32')
+            # flatten before checking shapes to avoid shape miss match
+            label = label.flat
+            pred_label = pred_label.flat
 
             labels, preds = check_label_shapes(label, pred_label)
 
-            self.sum_metric += (pred_label.flat == label.flat).sum()
-            self.num_inst += len(pred_label.flat)
+            self.sum_metric += (pred_label == label).sum()
+            self.num_inst += len(pred_label)
 
 
 @register