Thanks for the quick response.
With regards to your first suggestion, I’m a bit confused as my pos variable should be changing at every time step and at every step update. Thus, the gradient of pos with respect to energy should be changing as well. Using pos.grad after only calling the backwards step once doesn’t update the gradient at every step.
Memo-izing the gradients should help speed up the gradient calculation when I’m using a MH step in HMC, and I can try that. However, in this specific example, it may not help as my position shouldn’t ever be the same (I’m effectively drawing random samples from the energy probability distribution).