Thank you saan77, but I am still unable to understand how it is possible to get the Jacobian with a single backward pass.
The grad_tensors argument of backward 84 seems to work as a weighting mask for the tensors argument in the thread you posted.
I don’t want the gradient of my tensors to be accumulated at leaf nodes. I want to get the gradient of each of my tensors with respect to leaf nodes. Let’s say I have an image classifier, whose input has shape (batchsize, c, h, w) and its output has shape (batchsize, n_classes), I want the jacobian to have shape (batchsize, c, h, w, n_classes).
Did you manage to get something similar?