I think you are running into the classic problem with reverse mode automatic differentiation, where given a single output you can simultaneously compute the gradients of all inputs, but to take the gradients of multiple outputs, you must run backwards multiple times.
PyTorch doesn’t natively support forwards-mode AD, but there’s a cute trick you can use to get it: see https://github.com/pytorch/pytorch/issues/10223 43 (unfortunately, it is not very efficient)