diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 9bd3598b19..33fb36e8b6 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -561,7 +561,7 @@ def prepare_model_for_kbit_training(
             model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
 
         # Enable grads on non language models as well
-        requires_grad_for_gradient_checkpointing()
+        requires_grad_for_gradient_checkpointing(model)
     pass
 
     return model