Hello! I have this (simplified version) code:
grads = {}
def save_grad(name):
    def hook(grad):
        grads[name] = grad
    return hook

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(2, 1,  bias=False)
        self.linear2 = nn.Linear(1, 2,  bias=False)
            
    def forward(self, x):
        z = self.linear1(x)
        y_pred = self.linear2(z)

        return y_pred, z

for epoch in range(1000):
    model.train()
    for i, dt in enumerate(data.trn_dl):
        optimizer.zero_grad()
        output = model(dt[0])
        
        output[1].register_hook(save_grad('z_x_hat'))
        output[0][0].backward(retain_graph=True)
        z_x_hat = grads['z_x_hat']
        
        output[1].register_hook(save_grad('z_y_hat'))
        output[0][1].backward(retain_graph=True)
        z_y_hat = grads['z_y_hat']
        
        z_x_hat.requires_grad = True
        z_y_hat.requires_grad = True
                                     
        loss = abs(torch.sqrt(z_x_hat**2+z_y_hat**2)-1)
        print(z_x_hat,z_y_hat,loss)
        loss.backward()
        
        optimizer.step()

So my network has 2 outputs (x_hat, y_hat) and I want the loss function to be given by the derivative of the outputs with respect to the value of z (I need this as part of something more complex for the actual project, but this is the part where I am stuck). In this simple case, the derivatives are just the value of the weights from z to x_hat and z to y_hat (in the real case is the product of all the partial derivatives from x_hat to z). Using the code above:
output[1].register_hook(save_grad('z_x_hat'))
output[0][0].backward(retain_graph=True)
z_x_hat = grads['z_x_hat']

I am able to get the value of this derivative (and hence the weight), but the problem is that I need to use retain_graph=True in order to do that, and that screws up my whole backprop, as I want to optimize the network using backprop only on the loss at the end, not on these intermediate calls. Can someone help me with this? Thank you!