RuntimeError: Expected to mark a variable ready only once - with .backward() in validation_step

[kampelmuehler]

Hi all,
I'm getting RuntimeError: Expected to mark a variable ready only once when running a model on more than one GPU (this is my first time using lightning with DDP).

In my validation_step I need to loop through the model multiple times and run .backward() to get gradients - and this is what causes the error. The code looks something like this:

for _ in range(num_steps):
    # pred = model predictions
    # samples = input points to the network
    pred.sum().backward()

    gradient_norm = F.normalize(samples.grad.detach(), dim=2)
    samples = samples.detach()
    pred = pred.detach()
    samples = samples - gradient_norm * pred.reshape(-1, 1)
    samples = samples.detach()
    samples.requires_grad = True

for this purpose I also set

    def on_validation_model_eval(self, *args, **kwargs):
        super().on_validation_model_eval(*args, **kwargs)
        torch.set_grad_enabled(True)

It's exactly at pred.sum().backward() where the error is triggered.

I've tried find_unused_parameters=False and a static graph is impossible - so the options the error message gives are exhausted. As mentioned earlier it works fine on a single GPU.

Any hints on what causes that behavior and how it can be fixed?

I would also be fine with running validation on a single GPU, but I found that being impossible within lightning.