How to create class activation map during test_step?

[kapsner]

Hello,

I am trying to create a gradient class activation map using monai.visualize.GradCAM on images during test_step in order to support interpretability of predicted classes.

However, I constantly get the following error and assume that it is somehow related to pytorch_lightning's model states, but I do neither understand the source of this error nor how to fix/ work around it.

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

I have also also opened this discussion with more detail in the MONAI channel.

A minimal working example is provided here.

Would be great if anyone has an idea how to solve this.

Note: the error doesn't occur when running GradCAM from inside train_step