Update app.py

app.py

@@ -82,8 +82,11 @@ def gradio_callback(view_grad_cam, num_gradcam_images, view_misclassified, num_m
                     input_img, transparency = 0.5, target_layer_number = -1):
     
     confidence = inference_confidences(input_img, transparency = 0.5, target_layer_number = -1)
-    visualization = inference_visualization(input_img, transparency = 0.5, target_layer_number = -1)
-    return confidence, visualization
+    if view_grad_cam:
+        visualization = inference_visualization(input_img, transparency = 0.5, target_layer_number = -1)
+        return confidence, visualization
+    else:
+        return confidence
 
 
 title = "CIFAR10 trained on ResNet18 Model with GradCAM"