ResNet13 3D for Video classification GradCam

[GoncaloGuedes]

Hello, I have a 3D ResNet for video classification and I would like to use gradcam.
My input for the model are 16 RBG images with a size of 512. The video variable has a shape of [1, 3, 16, 512, 512]
My problem is when I run this code (like the example by the author)

targets = [ClassifierOutputTarget(1)]
target_layers = [model.model.layer4[-1]]
cam = GradCAM(model=model.model, target_layers=target_layers)
grayscale_cam = cam(input_tensor=video, targets=targets)

I get a grayscale_cam with a shape of [1, 512, 512, 32] instead of [1, 512, 512].

How can I fix it?
Thanks!!