From 95f872dedc2d6147f7314397f7b73db9fd5c730d Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 4 Mar 2025 19:20:38 -0800
Subject: [PATCH] Update llama.py

---
 unsloth/models/llama.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index fd1f16e30b..fa37fc34bf 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -2657,9 +2657,9 @@ class FastLlamaModel:
         pass
         m = model
         while hasattr(m, "model"):
-            _for_inference(m)
+            _for_training(m)
             m = m.model
-        _for_inference(m)
+        _for_training(m)
 
         # Also re-enable training for embeddings for NEFTune
         if hasattr(model, "get_input_embeddings"):