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"):