I don’t think that this is currently possible in geberal (unless iterating over the scalar parameters). The fundamental reason is that backpropagation in PyTorch isn’t prepared take deruvatives of vector-valued functions, so you are limited to taking the derivative of a scalar sum of derivatives, i.o.w. a Hessian-Vector product.
For a small number of parameters using torch.autograd.grad will help, but I’m not sure it scales to all parameters of large nets.
Best regards
Thomas