diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py
index d82c3d517f..c598c4cfe9 100644
--- a/deepmd/pt/utils/dataset.py
+++ b/deepmd/pt/utils/dataset.py
@@ -71,7 +71,7 @@ def get_frame_index(self):
                 row_indices_global = np.where(frames_with_elem)[0] + base_offset
                 element_counts[elem]["frames"] += len(row_indices)
                 element_counts[elem]["indices"].extend(row_indices_global.tolist())
-            base_offset += element_data.shape[0] 
+            base_offset += element_data.shape[0]
         element_counts = dict(element_counts)
         return element_counts