Minor fix

app.py

@@ -1,9 +1,9 @@
 import torch
 import dataset
-import visualize
 import albumentations
 from utils import get_misclassified_data
 from albumentations.pytorch import ToTensorV2
+from visualize import display_cifar_misclassified_data
 from torchvision import transforms
 import numpy as np
 import gradio as gr
@@ -88,7 +88,7 @@ def inference(input_img, is_grad_cam=True, transparency = 0.5, target_layer_numb
         # Get the misclassified data from test dataset
         misclassified_data = get_misclassified_data(model, device, test_loader)
         # Plot the misclassified data
-        visualize.display_cifar_misclassified_data(misclassified_data, number_of_samples=num_missclassified_images)
+        missclassified_images = display_cifar_misclassified_data(misclassified_data, number_of_samples=num_missclassified_images)
     else:
         missclassified_images = None
     

visualize.py

@@ -64,6 +64,8 @@ def display_cifar_misclassified_data(data: list,
         plt.xticks([])
         plt.yticks([])
 
+    return fig
+
 def display_gradcam_output(data: list,
                            model,
                            target_layers,