Inconsistent Device Error

[stefanbringuier]

When trying to invoke

...
aseAtoms.set_calculator = ORBCalculator(orbff, device='cuda')
aseAtoms.get_potential_energy()

I'm getting this issue with the PyTorch tensors:

  File "/home/user/orb-models/lib/python3.10/site-packages/ase/atoms.py", line 731, in get_potential_energy
    energy = self._calc.get_potential_energy(self)
  File "/home/user/orb-models/lib/python3.10/site-packages/ase/calculators/calculator.py", line 709, in get_potential_energy
    energy = self.get_property('energy', atoms)
  File "/home/user/orb-models/lib/python3.10/site-packages/ase/calculators/calculator.py", line 737, in get_property
    self.calculate(atoms, [name], system_changes)
  File "/home/user/orb-models/lib/python3.10/site-packages/orb_models/forcefield/calculator.py", line 67, in calculate
    out = self.model.predict(batch)
  File "/home/user/orb-models/lib/python3.10/site-packages/orb_models/forcefield/graph_regressor.py", line 492, in predict
    batch = self.model(batch)
  File "/home/user/orb-models/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/orb-models/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/orb-models/lib/python3.10/site-packages/orb_models/forcefield/gns.py", line 421, in forward
    batch = self.featurize_edges(batch)
  File "/home/user/orb-models/lib/python3.10/site-packages/orb_models/forcefield/gns.py", line 456, in featurize_edges
    rbfs = self.rbf(lengths)
  File "/home/user/orb-models/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/orb-models/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/orb-models/lib/python3.10/site-packages/orb_models/forcefield/rbf.py", line 97, in forward
    * (torch.exp(self.alpha * (-dist + self.cutoff_lower)) - self.means) ** 2
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

The only way I got this to work is if I do:

import torch
torch.set_default_device('cuda')
...
aseAtoms.set_calculator = ORBCalculator(orbff, device='cuda')
aseAtoms.get_potential_energy()

[CompRhys]

I also encountered this and was alternatively able to get it to work with:

orbff = pretrained.orb_v1()
orbff.to(device)
calc = ORBCalculator(orbff, device=device)

[benrhodes26]

Thanks for raising this and proposing fixes. We just released V0.3.2 which fixes this behaviour so GPU should work out-of-the-box.

We also added custom finetuning :)