Need an option for evaluation not in the inference_mode.

[LouChao98]

🚀 Feature

An option for the trainer run evaluation not in the inference_mode.

Motivation

In some scenerios, gradients are used even when evaluating the model. The recent PR (#12715) adds supports for the new PyTorch API inference_mode and enables it by default. However, unlike torch.no_grad that one can simply set torch.enable_grad to overwrite it locally, inference_mode do not provide a convenient way to quit it locally. This is because all tensors are marked and forbidden for furthur gradient-related usage. People have to clone tensors to make code working. Therefore, I suggest that pytorch-lightning provide a way to let users to choose whether to use the inference mode.

Pitch

Add an option to the Trainer class that allows users to decide whether to use the inference mode or not.

Alternatives

Additional context

From the pytorch's document https://pytorch.org/cppdocs/notes/inference_mode.html:

Inference tensors:
...
are immutable outside InferenceMode. So an error will be raised if you try to: - mutate their data outside InferenceMode. - mutate them into requires_grad=True outside InferenceMode. To work around you can make a clone outside InferenceMode to get a normal tensor before mutating.

cc @justusschock @kaushikb11 @awaelchli @Borda @rohitgr7 @akihironitta

[rohitgr7]

does this work?

def validation_step(self, ...):
    with torch.inference_mode(mode=False):
        ...

[LouChao98]

No, it does not work. The main issue is that we need to copy input tensors because tensors returned by the dataloader are all "Inference tensors". A workaround can be like this

def validation_step(self, batch):
    batch = {key: value.clone() for key, value in batch.items()}
    with torch.inference_mode(mode=False):

(Besides, this workaround is somewhat ugly.

[carmocca]

I think this is a fair ask. Also related: #13948