oops

mosaicml · Jul 25, 2024 · 244f8e1 · 244f8e1
1 parent b7c44ef
commit 244f8e1
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -465,8 +465,7 @@ def tensor_hook(
 
         hooks = []
         for _, module in state_dict_model.named_modules():
-            if isinstance(module, FSDP):
-                hooks.append(module._register_state_dict_hook(tensor_hook),)
+            hooks.append(module._register_state_dict_hook(tensor_hook),)
 
         state_dict = get_model_state_dict(
             state_dict_model,