How to use the S2EF model as ASE calcualtor with CUDA (GPU) about ocp HOT 1 CLOSED

Hi - Because ASE operations and methods are CPU-only, a CUDA calculator will have transferring overhead from CPU->GPU->CPU. With that said, depending on the application this overhead may be trivial compared to the model forward pass. I have opened a PR with these changes: #333. To create a CUDA supported calculator:

calc = OCPCalculator(checkpoint=/path/to/checkpoint.pt, device=0)
atoms.set_calculator(calc)
energy = atoms.get_potential_energy()

Where device is the GPU you wish to use if on a multi-gpu machine. If not specified, it'll default to the original CPU implementation. Hope this helps!

Note - To avoid the overhead you can always implement the ASE functions in pytorch to properly support GPU usage. This is what we did for the system relaxations, but I'd only recommend this if you plan on using a particular function or operation in a high throughput manner.

from ocp.