Update _utils.py

unslothai · Apr 9, 2024 · 35dbef8 · 35dbef8
1 parent d95daa0
commit 35dbef8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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)