using autograd in neural network calculation breaks the validation step
🐛 Bug
Hi,
just a small bug report, maybe this can be a feature in the future.
Assuming you want to have a derivative in your neural network as output,
the option: torch.set_grad_enabled(False), which is activated during validation, breaks the PyTorch lightning module.
So if you want to train and validate a model like this (see Code) in pytorch ligthning, a workaround is to  torch.set_grad_enabled(True) at the beginning of the validation step.
import torch 

class Feedforward(torch.nn.Module):
        def __init__(self, input_size, hidden_size):
            super(Feedforward, self).__init__()
            self.input_size = input_size
            self.hidden_size  = hidden_size
            self.fc1 = torch.nn.Linear(self.input_size, self.hidden_size)
            self.relu = torch.nn.ReLU()
            self.fc2 = torch.nn.Linear(self.hidden_size, 1)
     
        
        def forward(self, x):
            hidden = self.fc1(x)
            relu = self.relu(hidden)
            output = self.fc2(relu)
            output = output.sum()
            output =torch.autograd.grad(outputs=output, inputs=x, retain_graph=True, create_graph=True)
            return output[0]

test_input = torch.rand((10,3),requires_grad=True)
test_output = torch.rand((10,3))


model = Feedforward(3,10)
optim = torch.optim.Adam(model.parameters())
optim.zero_grad()
loss_fn = torch.nn.L1Loss()
model.train()
out = model(test_input)
loss = loss_fn(out, test_output)
loss.backward()
optim.step() 


I dont know if this behavior is intendet and if you want to find a workaround, but I leave this here as a Feature Request/ Bug Report.