fix: cuda

deepmodeling · Feb 28, 2024 · df33191 · df33191
1 parent c0f2c56
commit df33191
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/source/tests/pt/model/test_descriptor_se_r.py b/source/tests/pt/model/test_descriptor_se_r.py
@@ -137,8 +137,8 @@ def test_load_stat(self):
                         "r0": None,
                         "coord": torch.from_numpy(self.coord_ext).reshape(
                             -1, self.nall, 3
-                        ),
-                        "atype": torch.from_numpy(self.atype_ext),
+                        ).to(env.DEVICE),
+                        "atype": torch.from_numpy(self.atype_ext).to(env.DEVICE),
                         "box": None,
                         "natoms": self.nall,
                     }