diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py
index 495630a932..ae4a484909 100644
--- a/deepmd/utils/data.py
+++ b/deepmd/utils/data.py
@@ -103,13 +103,10 @@ def __init__(
                     f"Elements {missing_elements} are not present in the provided `type_map`."
                 )
             if not self.mixed_type:
-                # Use vectorized operation for better performance with large atom counts
-                # Create a mapping array where old_type_idx -> new_type_idx
-                max_old_type = max(self.atom_type) + 1
-                type_mapping = np.zeros(max_old_type, dtype=np.int32)
-                for old_idx in range(len(self.type_map)):
-                    type_mapping[old_idx] = type_map.index(self.type_map[old_idx])
-                self.atom_type = type_mapping[self.atom_type].astype(np.int32)
+                old_to_new_type_idx = np.array(
+                    [type_map.index(name) for name in self.type_map], dtype=np.int32
+                )
+                self.atom_type = old_to_new_type_idx[self.atom_type].astype(np.int32)
             else:
                 self.enforce_type_map = True
                 sorter = np.argsort(type_map)