diff --git a/hyperlora/model_loading.py b/hyperlora/model_loading.py
index d19ab42..0f7da38 100644
--- a/hyperlora/model_loading.py
+++ b/hyperlora/model_loading.py
@@ -138,6 +138,7 @@ def get_model(
         model = AutoModelForCausalLM.from_pretrained(**model_init_kwargs)
     else:
         model = MllamaForConditionalGeneration.from_pretrained(**model_init_kwargs)
+        model = model.language_model
     if peft_config is not None:
         model = PeftModel(model, peft_config)
     model.train(train)