So one way I could overcome this is to define the second order derivative as its own Pytorch module without using autograd, and then run my inputs through that in batches, which would be significantly faster because it would take advantage of the parallelism of batch processing as opposed to iterating over each element in a for loop. Of course, that’s kind of tedious to do, even with a standard MLP, especially if I want to use dropout. But it’s doable. Would it be possible to have batch vectorization in a future release? I think JAX autograd does it. In general, it would be very helpful to be able to treat derivatives of modules as modules. Then not only could you process in batches, but you could also reuse the same ones without having to recompute them. In functions where different inputs might have the same derivative function, you could also get away with reusing the same derivative module as opposed to recomputing a new one. Does that make sense?