I have seen JAX-MD[1] but not sure about "much better". On the other hand, there is just no MD implemented with PyTorch.

[1]: https://github.com/jax-md/jax-md