Place atomic numbers on same device.

chemle · Nov 29, 2024 · 10e90d6 · 10e90d6
1 parent 1268c7f
commit 10e90d6
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/emle/models/_ani.py b/emle/models/_ani.py
@@ -223,7 +223,9 @@ def __init__(
             # Optimise the ANI2x model if atomic_numbers are specified.
             if _has_nnpops and atomic_numbers is not None:
                 try:
-                    atomic_numbers = atomic_numbers.reshape(1, *atomic_numbers.shape)
+                    atomic_numbers = atomic_numbers.reshape(
+                        1, *atomic_numbers.shape
+                    ).to(self._device)
                     self._ani2x = _NNPOps.OptimizedTorchANI(
                         self._ani2x, atomic_numbers
                     ).to(device)