Unfortunately, the only thing autograd can compute is jacobian vector products. So if you want to reconstruct a full jacobian, you have to do backward with tensors [1, 0, 0] then [0, 1, 0] and [0, 0, 1].
You can look in the forum actually, there are a few posts discussing how to compute full Jacobian that might help you speed things up a bit