Update prediction.py

prediction.py

@@ -301,14 +301,14 @@ class Prediction:
         input_tensor = torch.stack(frames).unsqueeze(0)
         input_tensor = input_tensor.view(1, target_seq_length, 3, *self.resolution)
         input_tensor = input_tensor.to(self.device)
-        input_tensor.requires_grad_()
-    
-        # Forward pass
-        with torch.no_grad():
-            fmap, attn_wts, logits = self.model(input_tensor)
-    
+        
+        # Remove the torch.no_grad() context to allow gradient computation
+        input_tensor.requires_grad_(True)
+        
+        # Forward pass with gradient tracking enabled
+        fmap, attn_wts, logits = self.model(input_tensor)
+        
         # Register hook for Grad-CAM
-        fmap.requires_grad_()
         fmap.register_hook(self.save_gradients)
     
         # Get predictions for all classes
@@ -358,4 +358,3 @@ class Prediction:
             gradcam_image = None
     
         return prediction_string, gradcam_image, classification_details
-