fix map_location

huggingface · Nov 22, 2024 · 62fa626 · 62fa626
1 parent 26ac3e0
commit 62fa626
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/nanotron/serialize/optimizer.py b/src/nanotron/serialize/optimizer.py
@@ -351,7 +351,7 @@ def get_checkpoint_state_metadata(param_name: str, pp_rank: int, tp_rank: int) -
                     )
                     state_dict["state"][param_index][state_name] = sliced_tensor
 
-    optimizer.load_state_dict(state_dict, map_location="cpu")
+    optimizer.load_state_dict(state_dict, map_location=map_location)
 
 
 def load_lr_scheduler(