update

iterative · Apr 27, 2022 · 281e8ef · 281e8ef
1 parent a24f437
commit 281e8ef
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/demo/train.py b/demo/train.py
@@ -90,7 +90,7 @@ def get_confusion_image(predictions, dataset):
     for n, pred in enumerate(predictions):
         actual = pred["actual"]
         predicted = pred["predicted"]
-        image = np.array(dataset[n]) / 255
+        image = np.array(dataset[n][0]) / 255
         confusion[(actual, predicted)] = image
 
     max_i, max_j = 0, 0