Make sure to move data to CPU before calling vesin

[Luthaf]

Summary

Fix #155

I don't have access to a CUDA GPU, so I would appreciate either @frostedoyster or @orionarcher checking that this does work in practice!

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format.
  Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.
• All linting and tests pass.

Note that the CI system will run all the above checks. But it will be much more
efficient if you already fix most errors prior to submitting the PR. It is highly
recommended that you use the pre-commit hook provided in the repository. Simply run
pre-commit install and a check will be run prior to allowing commits.

[orionarcher]

Works! Thanks for the quick fix.

Question though, is there really no way to run vesin on GPU? I thought that was the role of vesin-torch.

[Luthaf]

Question though, is there really no way to run vesin on GPU? I thought that was the role of vesin-torch.

Not currently no =) We are working to create a CUDA and/or HIP implementation to make sure the code can run on GPU, hopefully in the next couple of months.

If you want a pure torch implementation of neighbor list, that can run on GPU (because it only uses functions from torch), you can try using https://github.com/felixmusil/torch_nl/, but I found that it is slower than vesin even when running on GPU, except when dealing with very large cells and large (>12A) cutoffs. https://github.com/openmm/NNPOps/ does contain some custom CUDA code for neighbor list, but requires a cell larger than 2 cutoff.

vesin-torch is only exposing the code from vesin in a TorchScript compatible way, and integrating with autograd.

[orionarcher]

Thanks for the explanation!

Will merge for now and look forward to the upcoming implementation!

[abhijeetgangan]

The torch version of neighbor lists which run on GPU are already included in torchsim including an API for torch_nl which accepts the batch as input. https://github.com/Radical-AI/torch-sim/blob/main/examples/scripts/7_Others/7.3_Batched_neighbor_list.py

[orionarcher]

If that could be subbed in it would be a nice addition to the metatensor model interface!

[Luthaf]

I'm not sure it is actually worth doing, since in all my tests torch_nl is slower than vesin, even when running on GPU. So unless someone has a system where this is not the case, I would rather keep the code simple!