style III

mosaicml · Dec 17, 2024 · 4301ceb · 4301ceb
1 parent 42c715a
commit 4301ceb
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -619,8 +619,7 @@ def tensor_hook(
 
         hooks = []
         for _, module in state_dict_model.named_modules():
-            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,