GradCAM on SegResNet leads to OOM

[vDawgg]

I am currently trying to get Class Activation Maps for a SegResNet model I have trained, which is pretty much defined in the same way as in the tutorial . The final model works quite well for the task that I have applied it to, but I cannot get the MONAI GradCAM implementation to work with SegResNet.

I have tried multiple different layers (conv_final, conv_final.2.conv, up_layers.2, etc.) but have in all cases gotten the following error:

• OutOfMemoryError: CUDA out of memory. Tried to allocate 6400.00 GiB.

This seems like something is completely misconfigured on my end, but Im a bit stumped on what to try next.

The only other error I have gotten was triggered, when I tried using 'up_layers.2.0.conv2.conv' as my target layer:

• RuntimeError: Output 0 of BackwardHookFunctionBackward is a view and is being modified inplace. This view was created inside a custom Function (or because an input was returned as-is) and the autograd logic to handle view+inplace would override the custom backward associated with the custom Function, leading to incorrect gradients. This behavior is forbidden. You can fix this by cloning the output of the custom Function.

Does this mean, that the way SegResNet is currently implemented makes it incompatible with the GradCAM implementation, or am I still missing something?

P.S.: The above asserts are also repro with the following snippet (depending on the chosen target layer):

import os
from monai.networks.nets import SegResNet
from monai.visualize.class_activation_maps import GradCAM
import torch

layer = 'up_layers.2'

os.environ['CUDA_VISIBLE_DEVICES'] = "0"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = SegResNet(
    blocks_down=[1, 2, 2, 4],
    blocks_up=[1, 1, 1],
    init_filters=16,
    in_channels=4,
    out_channels=3,
    dropout_prob=0.2,
).to(device)

cam = GradCAM(nn_module=model, target_layers=layer)
result = cam(x=torch.rand(1, 4, 128, 128, 80).to(device))