Merge pull request #149 from huggingface/fix-save-ckpt

[BUG] fix arg for save_checkpoint
huggingface · Apr 23, 2024 · b792a22 · b792a22
2 parents dfdc063 + 37cfcda
commit b792a22
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/nanotron/trainer.py b/src/nanotron/trainer.py
@@ -398,7 +398,7 @@ def train(
         self.pre_training(**kwargs)
 
         if self.config.checkpoints.save_initial_state and self.init_checkpoint_path is None:
-            self.save_checkpoint(dataloader_or_dls)
+            self.save_checkpoint()
 
         self.pipeline_engine: PipelineEngine = self.config.parallelism.pp_engine