Correct to_categorical for unsqueezed labels (closes #7)

(cherry picked from commit 7ba14c4) (cherry picked from commit 001b846)
Trusted-AI · Jun 19, 2018 · 266fc8b · 266fc8b
1 parent cf023bd
commit 266fc8b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/art/utils.py b/art/utils.py
@@ -41,7 +41,7 @@ def to_categorical(labels, nb_classes=None):
     if not nb_classes:
         nb_classes = np.max(labels) + 1
     categorical = np.zeros((labels.shape[0], nb_classes), dtype=np.float32)
-    categorical[np.arange(labels.shape[0]), labels] = 1
+    categorical[np.arange(labels.shape[0]), np.squeeze(labels)] = 1
     return categorical