Computing batch Jacobian efficiently

Not sure what you mean by “one input point” could you clarify?

For nn.Module, you can check this answer: Get gradient and Jacobian wrt the parameters - #3 by albanD