I\'m writing some code to simulate the motion of a system of particles on GPU - similar to N-body system. I have a version in numpy and a direct port of that to pytorch, wh