Update Phi-3 train.py to fix model import

mbzuai-oryx · Apr 30, 2024 · c95a8e1 · c95a8e1
1 parent 31c9582
commit c95a8e1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/Phi-3-V/train.py b/Phi-3-V/train.py
@@ -923,7 +923,7 @@ def train(attn_implementation=None):
  **bnb_model_from_pretrained_args
  )
  else:
- model = LlavaLlamaForCausalLM.from_pretrained(
+ model = LlavaPhiForCausalLM.from_pretrained(
  model_args.model_name_or_path,
  cache_dir=training_args.cache_dir,
  attn_implementation=attn_implementation,