fix

mosaicml · Dec 17, 2024 · 3a7c0f1 · 3a7c0f1
1 parent 9103dbd
commit 3a7c0f1
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -788,6 +788,7 @@ def _save_checkpoint(
         dist.barrier()
 
         if dist.get_global_rank() == 0:
+            assert new_model_instance is not None
             if upload_to_save_folder:
                 # This context manager casts the TE extra state in io.BytesIO format to tensor format
                 # Needed for proper hf ckpt saving.