As part of my model, I want to extract the gradients of my model output relative to its input. Such that my model looks something like this:
class Model(nn.Module):
        .
        .
        .
       def forward(self, input_data):
              output = MLP(input_data) 
              gradients = grad(output,input_data,grad_outputs=torch.ones_like(output))
              return output, gradients

Typically, my training loop looks as such:
 for epoch in range(num_epochs):

        for data_sample in dataloader:
            input_data=data_sample[0]
            target=data_sample[1]
            batch_size=len(target)
            target=target.reshape(400,1)

            input_data=input_data.to(device)
            target=target.to(device)

            def closure():
                optimizer.zero_grad()
                output, gradients=model(input_data)
                loss=criterion(output,target)
                loss.backward()
                return loss

            loss=optimizer.step(closure)


The idea is to then implement the gradients as part of the loss function for my given application. My issue is that when run: I get a “Trying to backward through the graph a second time, but the buffers have already been freed” error. When suppressing the line of code containing grad() in my model, training works normally. I read something along the lines that gradients are freed when .backward() is called, but not quite sure how to get around this. I tried changing the retain_graph argument in grad() to True but then I get an error saying "element 0 of tensors does not require grad and does not have a grad_fn). Any help is appreciated…Thanks!