Hi,
I’m facing an issue when trying to use the jacobian of a computation as part of the loss function.
I created this minimal example with two networks.
import torch
torch.manual_seed(0)

x=torch.rand(1,3)
net1=torch.nn.Sequential( torch.nn.Linear(3,2), torch.nn.Tanh() )
net2=torch.nn.Sequential( torch.nn.Linear(2,1), torch.nn.Tanh() )


hidden=net1(x)
output=net2(hidden)

jac=torch.autograd.grad(output, hidden, grad_outputs=torch.ones_like(output), create_graph=True, retain_graph=True)[0]

loss=jac.norm()
loss.backward()

print("net1 grad", net1[0].weight.grad.norm()) #this should be non zero
print("net2 grad", net2[0].weight.grad.norm()) #But how to make this zero?

The jacobian of the output with respect to the hidden is used in the loss function
After calling loss.backward() gradients of net1 and net2 are populated (as expected)
However, I would like to optimize only net1 but keep the weights of net2 frozen.
From my understanding, the gradient flows backwards through torch.autograd.grad and through output and finally into net2
I tried having
output.register_hook(lambda grad: torch.zeros_like(grad))
loss.backward()

which from my understaning should set the gradient flowing back from output to zero and therefore net2 should essentially frozen, but this is not the case.
Do you have any suggestions on how to solve this?
Thank you very much for your help!