diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 663cd0338a..88940fcbdc 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -245,7 +245,10 @@ 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))], device=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