pytorch-grad-cam: AxisError: axis 2 is out of bounds for array of dimension 2

Getting the following error when trying out the cam function on an image example. This might be an issue with how I have loaded in my data, but not sure how to debug it.

Code to reproduce:

from pytorch_grad_cam import GradCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
import PIL


target_layers = [model.linear_layers[-1]]
img, label, path = next(iter(test_loader))
img, label = img.to(DEVICE), label.to(DEVICE)

img = img.float()

cam = GradCAM(model=model, target_layers=target_layers)

target_category = None

grayscale_cam = cam(input_tensor=img)

grayscale_cam = grayscale_cam[0, :]
visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)

Note the image is of shape torch.Size([64, 1, 128, 128]) Here is the traceback:

---------------------------------------------------------------------------
AxisError                                 Traceback (most recent call last)
<ipython-input-188-441d284cf4e0> in <module>()
     14 target_category = None
     15 
---> 16 grayscale_cam = cam(input_tensor=img)
     17 
     18 grayscale_cam = grayscale_cam[0, :]

7 frames
/usr/local/lib/python3.7/dist-packages/pytorch_grad_cam/base_cam.py in __call__(self, input_tensor, targets, aug_smooth, eigen_smooth)
    183 
    184         return self.forward(input_tensor,
--> 185                             targets, eigen_smooth)
    186 
    187     def __del__(self):

/usr/local/lib/python3.7/dist-packages/pytorch_grad_cam/base_cam.py in forward(self, input_tensor, targets, eigen_smooth)
     93         cam_per_layer = self.compute_cam_per_layer(input_tensor,
     94                                                    targets,
---> 95                                                    eigen_smooth)
     96         return self.aggregate_multi_layers(cam_per_layer)
     97 

/usr/local/lib/python3.7/dist-packages/pytorch_grad_cam/base_cam.py in compute_cam_per_layer(self, input_tensor, targets, eigen_smooth)
    128                                      layer_activations,
    129                                      layer_grads,
--> 130                                      eigen_smooth)
    131             cam = np.maximum(cam, 0)
    132             scaled = scale_cam_image(cam, target_size)

/usr/local/lib/python3.7/dist-packages/pytorch_grad_cam/base_cam.py in get_cam_image(self, input_tensor, target_layer, targets, activations, grads, eigen_smooth)
     52                                        targets,
     53                                        activations,
---> 54                                        grads)
     55         weighted_activations = weights[:, :, None, None] * activations
     56         if eigen_smooth:

/usr/local/lib/python3.7/dist-packages/pytorch_grad_cam/grad_cam.py in get_cam_weights(self, input_tensor, target_layer, target_category, activations, grads)
     20                         activations,
     21                         grads):
---> 22         return np.mean(grads, axis=(2, 3))

<__array_function__ internals> in mean(*args, **kwargs)

/usr/local/lib/python3.7/dist-packages/numpy/core/fromnumeric.py in mean(a, axis, dtype, out, keepdims)
   3371 
   3372     return _methods._mean(a, axis=axis, dtype=dtype,
-> 3373                           out=out, **kwargs)
   3374 
   3375 

/usr/local/lib/python3.7/dist-packages/numpy/core/_methods.py in _mean(a, axis, dtype, out, keepdims)
    145 
    146     is_float16_result = False
--> 147     rcount = _count_reduce_items(arr, axis)
    148     # Make this warning show up first
    149     if rcount == 0:

/usr/local/lib/python3.7/dist-packages/numpy/core/_methods.py in _count_reduce_items(arr, axis)
     64     items = 1
     65     for ax in axis:
---> 66         items *= arr.shape[mu.normalize_axis_index(ax, arr.ndim)]
     67     return items
     68 

AxisError: axis 2 is out of bounds for array of dimension 2