Update llama.py

2025-03-04 19:20:38 -08:00 · 2025-03-04 19:20:38 -08:00 · 95f872dedc
commit 95f872dedc
parent e887f43b52
1 changed files with 2 additions and 2 deletions
--- 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"):