fix: UTs

deepmodeling · Mar 5, 2024 · d436444 · d436444
1 parent a30bc35
commit d436444
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -245,9 +245,7 @@ def remap_atype(
         ), "The input `atype` cannot be handled by the type_map."
         type_2_idx = {atp: idx for idx, atp in enumerate(ori_map)}
         # this maps the atype in the new map to the original map
-        mapping = torch.tensor(
-            [type_2_idx[new_map[idx]] for idx in range(len(new_map))]
-        ).to(atype.device)
+        mapping = torch.tensor([type_2_idx[new_map[idx]] for idx in range(len(new_map))], device=atype.device)
         updated_atype = mapping[atype.long()]
         return updated_atype