Extract and speed up device transfers

[frostedoyster]

Moving systems to the device can be optimized as it requires moving many tensors sequentially (cell, positions, types, NLs, times the number of systems in a batch).

no compilation: 4.7 ms
script: 3.9 ms
torch.compile: 4.0 ms

---

📚 Documentation preview 📚: https://metatrain--357.org.readthedocs.build/en/357/

[frostedoyster]

@Luthaf do you know what's happening in the docs?

[Luthaf]

How does this speed it up? This is faster by using torchscript? Also the gains you show above are negligible, do you also have timings for a more realistic training set?

---

@Luthaf do you know what's happening in the docs?

This looks like trying to compile to torchscript the documentation classes, which only contain documentation. Maybe disabling TorchScript for the docs would help? There is an environment variable for this.

[frostedoyster]

The speed-up is nice, but I think the major advantage is to extract a function that would be essentially used everywhere

[frostedoyster]

We're already using the environment variable, right?

[Luthaf]

The docs are already using one env variable to switch between the TorchScript version of the classes and the documentation version of the classes. There is a separate env variable to disable TorchScript compilation altogether as well.