bug fix

app.py

@@ -55,6 +55,9 @@ def visualize_single_image(image_path):
 
     # MAE reconstruction pasted with visible patches
     im_paste = img * (1 - mask) + predicted_img * mask
+
+    # convert to numpy
+    img = img[0].cpu().detach().numpy()
     
 #    # make the plt figure larger
 #     plt.figure(figsize=(12, 4))
@@ -73,7 +76,7 @@ def visualize_single_image(image_path):
 
 #     plt.tight_layout()
 
-    return img[0].numpy(), im_masked[0].numpy(), predicted_img[0].numpy(), im_paste[0].numpy()
+    return img # , im_masked[0].numpy(), predicted_img[0].numpy(), im_paste[0].numpy()
 
 inputs_image = [
     gr.components.Image(type="filepath", label="Input Image"),