Thanks for your reply! I did consider using .detach but it doesn’t seem to fit my use case, as I need to backpropagate through the gradient operation. I ended up using torch.autograd.grad with retain_graph=True and create_graph=True and it appears to be doing the right thing.
I do have one follow-up question: in torch.autograd.grad, what is the grad_outputs argument used for? The description in the docs is unclear to me. I’m leaving it as the default value and my code appears to be working fine.