diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 77509cd8da..69e570dc9a 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -96,7 +96,7 @@ def prepare_model_for_kbit_training(
 
     # Freeze all parameters except LoRA
     import re
-    with torch.inference_mode():
+    with torch.no_grad():
         for name, param in model.named_parameters():
             if ".lora_A." in name or ".lora_B." in name or ".lora_magnitude_vector" in name:
                 param.requires_grad_(True)