update

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