Re-enabling gradients in Validation loop?

[harritaylor]

I'm trying to use Layer-wise relevance propagation as part of the training and validation loop of training my model, but it requires gradients to be present in order to calculate the relevance. I'm trying to figure out the best way of approaching this in a Lightning-friendly way. I can already do it for the training loop, but the validation loop is what's giving me problems. The only way I can think to do it at the moment is once the validation loop finishes, re-run the validation dataset with gradients enabled without calling the optimiser's .step() function. I'd really appreciate some guidance here before I dive in. Thanks!

[awaelchli]

You can call torch.set_grad_enabled(True) in your validation loop or in any hook you want. Doesn't that work?
And perhaps model.train() if you have normalization layers dropout or that kind of layers.

[harritaylor]

Fantastic, I didn't realise that was a feature. I have been searching documentation all morning trying to find something like that. I'll give it a shot. Thanks!