diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 0f0d4c159f..2423e8f942 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "2025.3.1"
+__version__ = "2025.3.4"
 
 __all__ = [
     "SUPPORTS_BFLOAT16",
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index bcabbd5125..a5bc8712e9 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -1538,6 +1538,7 @@ def _wrap_fast_inference(generate, device_type, dtype, model):
     # Wraps inference with bfloat16 / float16
     @torch.inference_mode
     def _fast_generate(*args, **kwargs):
+        if hasattr(model, "for_inference"): model.for_inference()
 
         if hasattr(model, "config") and hasattr(model.config, "max_position_embeddings"):
             if "input_ids" in kwargs and kwargs["input_ids"] is not None and "max_new_tokens" in kwargs:
@@ -1603,6 +1604,9 @@ def _wrap_fast_inference(generate, device_type, dtype, model):
             accelerate.utils.operations.send_to_device = accelerate_old_send_to_device
         pass
 
+        # Return to training state
+        if hasattr(model, "for_training"): model.for_training()
+
         return output
     pass
     return _fast_generate
@@ -2416,6 +2420,9 @@ class FastLlamaModel:
             model.load_lora = partial(load_lora, model)
         pass
 
+        # Add for_inference and for_training
+        model.for_training  = partial(FastLlamaModel.for_training,  model)
+        model.for_inference = partial(FastLlamaModel.for_inference, model)
         return model
     pass