From bd506f2c449ae9c0e6d480fd715d114a0cfceee0 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 1 Jan 2025 22:34:40 -0800
Subject: [PATCH] Update _utils.py

---
 unsloth/models/_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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