style II

mosaicml · Dec 17, 2024 · 42c715a · 42c715a
1 parent 44e609e
commit 42c715a
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -620,8 +620,7 @@ def tensor_hook(
         hooks = []
         for _, module in state_dict_model.named_modules():
             hooks.append(
-                module._register_state_dict_hook(tensor_hook),
-            )
+                module._register_state_dict_hook(tensor_hook),)
 
         state_dict = get_model_state_dict(
             state_dict_model,