Merge branch 'fix-assert-warning' into 'main'

Fix assert warning in !2282 See merge request ADLR/megatron-lm!2443
NVIDIA · Dec 12, 2024 · fd69c2f · fd69c2f
2 parents 371feef + 2816445
commit fd69c2f
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/megatron/core/rerun_state_machine.py b/megatron/core/rerun_state_machine.py
@@ -252,8 +252,7 @@ def train_step(data_iterator, ...):
             for d in data_iterators:
                 assert (
                     isinstance(d, RerunDataIterator),
-                    "data iterator is not wrapped with RerunDataIterator",
-                )
+                ), "data iterator is not wrapped with RerunDataIterator"
 
         # Are we about to start the initial run?
         if self.state == RerunState.NOT_RUNNING_YET:
@@ -263,8 +262,7 @@ def train_step(data_iterator, ...):
             if self.data_iterator_checkpoints is not None:
                 assert (
                     len(self.data_iterator_checkpoints) == len(data_iterators),
-                    "data_iterator has different length than checkpointed data iterator",
-                )
+                ), "data iterator has different length than checkpointed data iterator"
                 for i, d in enumerate(data_iterators):
                     d.set_checkpoint_state(self.data_iterator_checkpoints[i])
                 self.data_iterator_checkpoints = None
@@ -667,8 +665,7 @@ def save_my_model_checkpoint(data_iterator, ...):
         for d in data_iterators:
             assert (
                 isinstance(d, RerunDataIterator),
-                "data iterator is not wrapped with RerunDataIterator",
-            )
+            ), "data iterator is not wrapped with RerunDataIterator"
 
         state: dict[str, Any] = {
             'mode': self.mode,