fix

mosaicml · Jan 25, 2025 · 792c720 · 792c720
1 parent 6e2a3b8
commit 792c720
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -37,6 +37,7 @@
 )
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from transformers import (
+    AutoModel,
     PretrainedConfig,
     PreTrainedModel,
     PreTrainedTokenizerBase,
@@ -205,7 +206,7 @@ def save_model_patch(*args: Any, **kwargs: Any):
 
     if is_peft:
         transformers_in_memory_model = {
-            'model': transformers_model,
+            'model': AutoModel.from_pretrained(transformers_model),
             'tokenizer': tokenizer,
         }
     else: