Update prediction.py

prediction.py

@@ -303,14 +303,9 @@ class Prediction:
         input_tensor = input_tensor.to(self.device)
         input_tensor.requires_grad_()
     
-        # Forward pass to get model output dict
+        # Forward pass
         with torch.no_grad():
-            output_dict = self.model(input_tensor)
-    
-        # Extract relevant outputs
-        fmap = output_dict["fmap"]
-        attn_wts = output_dict["attn"]
-        logits = output_dict["logits"]
+            fmap, attn_wts, logits = self.model(input_tensor)
     
         # Register hook for Grad-CAM
         fmap.requires_grad_()