Compute gradients in parallel
The developers docs for torch specify that the arguments can be tensors, so you can just pass the tensors directly, but you need to reshape your output so that j is the last index. First, rewrite your code so that j is the last index and in your for-loop you have

l1 = torch.autograd.grad(output[0][0][j], ...)


And then drop the for-loop and remove the j-indices.