long -> int32

Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
njzjz · Sep 21, 2024 · 8c8553f · 8c8553f
1 parent 83302a1
commit 8c8553f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/source/tests/pt/common.py b/source/tests/pt/common.py
@@ -72,7 +72,7 @@ def eval_model(
         if spins is not None:
             assert isinstance(spins, torch.Tensor), err_msg
         assert isinstance(atom_types, torch.Tensor) or isinstance(atom_types, list)
-        atom_types = torch.tensor(atom_types, dtype=torch.long, device=DEVICE)
+        atom_types = torch.tensor(atom_types, dtype=torch.int32, device=DEVICE)
     elif isinstance(coords, np.ndarray):
         if cells is not None:
             assert isinstance(cells, np.ndarray), err_msg