train.py

icon-lab · Sep 24, 2024 · 8d4fa03 · 8d4fa03
1 parent 26f63cb
commit 8d4fa03
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/train.py b/train.py
@@ -169,23 +169,23 @@ def calculate_psnr(original, reconstructed):
                     img = img.data.squeeze().cpu()
 
                     ax = fig.add_subplot(spec[0])
-                    ax.imshow(img.permute(1, 2, 0), cmap="gray", vmin=0, vmax=1)
+                    ax.imshow(img.squeeze(), cmap="gray", vmin=0, vmax=1)
                     ax.axis("off")
                     ax.set_title("Input")
 
                     img = recon_train[0]
                     img = img.data.squeeze().cpu()
 
                     ax = fig.add_subplot(spec[1])
-                    ax.imshow(img.permute(1, 2, 0), cmap="gray", vmin=0, vmax=1)
+                    ax.imshow(img.squeeze(), cmap="gray", vmin=0, vmax=1)
                     ax.axis("off")
                     ax.set_title("Reconstructed")
 
                     img = target[0]
                     img = img.data.squeeze().cpu()
 
                     ax = fig.add_subplot(spec[2])
-                    ax.imshow(img.permute(1, 2, 0), cmap="gray", vmin=0, vmax=1)
+                    ax.imshow(img.squeeze(), cmap="gray", vmin=0, vmax=1)
                     ax.axis("off")
                     ax.set_title("Ground Truth")