fix index type

wanghan-iapcm · Oct 23, 2024 · a7427c9 · a7427c9
1 parent 1a1f145
commit a7427c9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
@@ -204,11 +204,9 @@ def forward_atomic(
             self.tab.ntypes, self.tab.ntypes, self.tab.nspline, 4
         )
 
-        # (nframes, nloc, nnei), dtype is the same as atype.
+        # (nframes, nloc, nnei), index type is int64.
         j_type = extended_atype[
-            np.arange(extended_atype.shape[0], dtype=extended_atype.dtype)[
-                :, None, None
-            ],
+            np.arange(extended_atype.shape[0], dtype=np.int64)[:, None, None],
             masked_nlist,
         ]
 

diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -272,7 +272,7 @@ def forward_atomic(
             torch.arange(
                 extended_atype.size(0),
                 device=extended_coord.device,
-                dtype=extended_atype.dtype,
+                dtype=torch.int64,
             )[:, None, None],
             masked_nlist,
         ]