fix: model mismatch

Author: jp1924

test/convergence/bf16/test_mini_models_multimodal.py

@@ -68,7 +68,6 @@
     MLLAMA_AVAILABLE = False
 
 try:
-
     from transformers import CLIPImageProcessor
     from transformers import CLIPVisionConfig
     from transformers import LlamaConfig
@@ -617,6 +616,8 @@ def run_mini_model_multimodal(
 
     set_seed(42)
 
+    model = create_model(model_name).to(dtype).to(device)
+
     revert_kwargs = {"model_config": MINI_MODEL_SETUPS[model_name]}
     if "mllama" in model_name:
         revert_kwargs["model_type"] = "conditional_generation"
@@ -636,13 +637,12 @@ def run_mini_model_multimodal(
         else:
             kwargs["swiglu"] = True
 
-        kwargs["model"] = create_model(model_name)
+        kwargs["model"] = model
 
         MINI_MODEL_SETUPS[model_name].liger_kernel_patch_func(**kwargs)
     else:
         MINI_MODEL_SETUPS[model_name].liger_kernel_patch_revert_func(**revert_kwargs)
 
-    model = create_model(model_name).to(dtype).to(device)
     model.gradient_checkpointing_enable()
 
     train_dataset = create_multimodal_dataset(model_name)