work

swiss-ai · Nov 26, 2024 · 36d5819 · 36d5819
1 parent c4effa9
commit 36d5819
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 368 deletions.
diff --git a/examples/custom-dataloader/README.md b/examples/custom-dataloader/README.md
diff --git a/examples/custom-dataloader/config_custom_dl.yaml b/examples/custom-dataloader/config_custom_dl.yaml
diff --git a/examples/custom-dataloader/run_train.py b/examples/custom-dataloader/run_train.py
diff --git a/src/nanotron/serialize/optimizer.py b/src/nanotron/serialize/optimizer.py
@@ -282,7 +282,7 @@ def get_checkpoint_state_metadata(param_name: str, pp_rank: int, tp_rank: int) -
         state_dict = torch.load(
             root_folder
             / optimizer_filename(parallel_context, is_zero=optimizer.inherit_from(optim.ZeroDistributedOptimizer)),
-            map_location=map_location,
+            map_location=map_location, weights_only=True
         )
 
     if isinstance(optimizer, ZeroDistributedOptimizer):
@@ -315,5 +315,5 @@ def load_lr_scheduler(
 ):
     root_folder = root_folder / "lr_scheduler"
 
-    state_dict = torch.load(root_folder / lr_scheduler_filename(parallel_context))
+    state_dict = torch.load(root_folder / lr_scheduler_filename(parallel_context), weights_only=True)
     lr_scheduler.load_state_dict(state_dict)