Autograd.grad computation time becomes longer and longer

Yes, I need to iterate this for NT times, and I need to keep all values of pos and vel as my generated data set…