bug fix to model predictions (#345)

microsoft · Jun 5, 2020 · 4986bec · 4986bec
1 parent 55bad56
commit 4986bec
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/experiments/interpretation/dutchf3_patch/local/test.py b/experiments/interpretation/dutchf3_patch/local/test.py
@@ -236,7 +236,7 @@ def _patch_label_2d(
                 # save image:
                 image_to_disk(np.array(batch[i, 0, :, :]), path_prefix + "_img.png")
                 # dump model prediction:
-                mask_to_disk(model_output[i, :, :, :].argmax(dim=1).numpy(), path_prefix + "_pred.png", num_classes)
+                mask_to_disk(model_output[i, :, :, :].argmax(dim=0).numpy(), path_prefix + "_pred.png", num_classes)
                 # dump model confidence values
                 for nclass in range(num_classes):
                     image_to_disk(