Upload 2 files

pytorch_grad_cam/base_cam.py

@@ -19,6 +19,8 @@ class BaseCAM:
         self.cuda = use_cuda
         if self.cuda:
             self.model = model.cuda()
+        else:
+            self.model = model.cpu()
         self.reshape_transform = reshape_transform
         self.compute_input_gradient = compute_input_gradient
         self.uses_gradients = uses_gradients
@@ -62,6 +64,7 @@ class BaseCAM:
             cam = get_2d_projection(weighted_activations)
         else:
             cam = weighted_activations.sum(axis=1)
+
         return cam
 
     def forward(self, input_tensor, target_encoding, target_category=None, eigen_smooth=False):
@@ -73,6 +76,7 @@ class BaseCAM:
                                                    requires_grad=True)
         # output will be the image encoding
         output = self.activations_and_grads(input_tensor)
+
         if isinstance(target_category, int):
             target_category = [target_category] * input_tensor.size(0)