🐛 Bug

AssertionError: Tensor target dimension torch.Size([1, 1, 224, 224]) is not valid. torch.Size([1, 1, 224, 224]) I am facing this error even though the target shape and output shape in the mentioned error is same . I am trying to get the GRAD CAM working on segmentation model which is mobilnet + unet .

`from captum.attr import LayerGradCam

Assuming train_set_raw is a dataset containing (image, mask) pairs

image, mask = train_set_raw[0] input_tensor = image.unsqueeze(0).to(DEVICE)

Assuming your model outputs segmentation masks

You need to provide a target mask for Grad-CAM

target_tensor = (mask > 0.5).float().unsqueeze(0).to(DEVICE) # Ensure correct dimensions

grad_cam = LayerGradCam(model, model.encoder.model.blocks[-1][0].conv) # Assuming you want to visualize the segmentation head

Calculate Grad-CAM attribution

attr = grad_cam.attribute(input_tensor, target=target_tensor, relu_attributions=False)

Convert the attribution to numpy and plot the heatmap

heatmap = attr.squeeze().cpu().numpy() plt.imshow(heatmap, cmap='viridis') plt.axis('off') plt.show()`

The error is faced on google collab T4 gpu .