Need to clean up the key some more

allenai · Dec 17, 2023 · 52bcd3f · 52bcd3f
1 parent a0d480b
commit 52bcd3f
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/olmo/checkpoint.py b/olmo/checkpoint.py
@@ -646,6 +646,7 @@ def restore_checkpoint(
                                 continue
                             key = f"{module_name}.{fqn}"
                             key = key.replace("_fsdp_wrapped_module.", "")
+                            key = key.lstrip(".")
                             t = state_dict_to_load[key]
                             t = t.flatten()
                             param[spi.offset_in_shard : spi.offset_in_shard + spi.numel_in_shard].copy_(