PyTorch loss function that depends on gradient of network with respect to input
I'm trying to implement a loss function that depends on the gradient of the network with respect to its inputs. That is, the loss function has a term like
sum(u - grad_x(network(x)))
where u is computed by forward propagating x through the network.
I'm able to compute the gradient by calling
funcApprox = funcNetwork.forward(X)
funcGrad = grad(funcApprox, X, grad_outputs=torch.ones_like(funcApprox))

Here, funcNetwork is my NN and X is the input. These computations are done in the loss function.
However, now if I attempt to do the following
opt.zero_grad()
loss = self.loss(X) # My custom loss function that calculates funcGrad, etc., from above
            
opt.zero_grad()
loss.backward()
opt.step()

I see the following error:
RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.

on the loss.backward() line from above.
I've tried playing around with create_graph, retain_graph, etc. but to no avail.
Any help is appreciated!

As per comment by @aretor, setting retain_graph=True, create_graph=False in the grad call in the loss function, and retain_graph=True in backward solves the issue.
