Computing batch Jacobian efficiently

[albanD]

Note that if you’re using the latest version of pytorch, there is a vectorize=True flag for functional.jacobian() that might speed things up in some cases