Merge pull request #258 from eliebak/quick-fix-save-lr

quick fix
huggingface · Dec 3, 2024 · 4cf36a8 · 4cf36a8
2 parents cfcdeae + 521f759
commit 4cf36a8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/nanotron/serialize/main.py b/src/nanotron/serialize/main.py
@@ -105,6 +105,7 @@ def save(
 
             save_lr_scheduler(
                 lr_scheduler=lr_scheduler,
+                is_zero=config.optimizer.zero_stage,
                 parallel_context=parallel_context,
                 root_folder=root_folder,
             )