minor bug fix!

app.py

@@ -58,11 +58,11 @@ def visualize_single_image_no_mode(image_path):
     # MAE reconstruction pasted with visible patches
     im_paste = img * (1 - mask) + predicted_img * mask
 
-    # resize the image to 96 x 96
-    img = v2.functional.resize(img[0], (96, 96))
-    im_masked = v2.functional.resize(im_masked[0], (96, 96))
-    predicted_img = v2.functional.resize(predicted_img[0], (96, 96))
-    im_paste = v2.functional.resize(im_paste[0], (96, 96))
+    # remove the batch dimension
+    img = img[0]
+    im_masked = im_masked[0]
+    predicted_img = predicted_img[0]
+    im_paste = im_paste[0]
     
     # make the plt figure larger
     plt.figure(figsize=(18, 8))
@@ -101,6 +101,7 @@ def visualize_single_image_pca_mode(image_path):
     im_paste = img * (1 - mask) + predicted_img * mask
 
     # remove the batch dimension
+    img = img[0]
     im_masked = im_masked[0]
     predicted_img = predicted_img[0]
     im_paste = im_paste[0]