Hi,
Yes this looks like the right way to do it.
FYI we now have a built in function that does the same thing: https://pytorch.org/docs/stable/autograd.html#torch.autograd.functional.jacobian 217
There is no better way to compute the jacobian yet I’m afraid. But we’re working on it.