fix

allenai · Apr 9, 2024 · 7312c71 · 7312c71
1 parent dc2a9d9
commit 7312c71
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/benchmarks/fsdp/test.py b/src/benchmarks/fsdp/test.py
@@ -60,7 +60,7 @@ def main(
     print_rank0(f"Loading OLMo-core FSDP checkpoint from {checkpoint_dir}...")
     load_model_and_optim_state(checkpoint_dir, olmo_model, olmo_optim)
 
-    print_rank0(f"Checking state dict...")
+    print_rank0("Checking state dict...")
     with TorchFSDP.summon_full_params(torch_model), olmo_model.summon_full_params():
         torch_state_dict = {k.replace("_fsdp_wrapped_module.", ""): v for k, v in torch_model.state_dict().items()}
         olmo_state_dict = olmo_model.state_dict()