diff --git a/src/nanotron/serialize/main.py b/src/nanotron/serialize/main.py
index e9ad04d8..b1445b48 100644
--- 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,
             )