Thanks for your response! What you’re saying sounds like it might work; I’ll give it a try.
In terms of why this leads to double counting - I’m calling autograd.grad with these tensors as the inputs, and the next thing that’s going to happen is that (essentially) .backward(calculated_gradient) will be called on each of those tensors.
(Actually this is all happening inside a autograd.Function.backward and the calculated gradients will be passed onto those tensors for the autograd framework to compute in the usual way. This is all part of some custom-backward memory saving recomputation shenanigans.)