From 020cbd2dd50204e46885cfa5b121805efedebb3b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 25 Sep 2024 00:22:54 -0700
Subject: [PATCH] Update _utils.py

---
 unsloth/models/_utils.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index f868c855bc..cd66825edf 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -41,6 +41,8 @@ __all__ = [
     "torch_amp_custom_bwd",
     "accelerate_old_send_to_device",
     "accelerate_new_send_to_device",
+    "patch_gradient_checkpointing",
+    "unpatch_gradient_checkpointing",
 ]
 
 import torch
@@ -791,7 +793,7 @@ class Unsloth_Offloaded_Gradient_Checkpointer(torch.autograd.Function):
     def backward(ctx, dY):
         (hidden_states,) = ctx.saved_tensors
         hidden_states = hidden_states.to("cuda:0", non_blocking = True).detach()
-        hidden_states.requires_grad = True
+        hidden_states.requires_grad_(True)
         with torch.enable_grad():
             (output,) = ctx.forward_function(hidden_states, *ctx.args)
         torch.autograd.backward(output, dY)
@@ -806,6 +808,17 @@ def unsloth_offloaded_gradient_checkpoint(function, *args, use_reentrant = None,
 pass
 
 
+import torch.utils
+old_checkpoint = torch.utils.checkpoint
+def patch_gradient_checkpointing():
+    torch.utils.checkpoint = unsloth_offloaded_gradient_checkpoint
+pass
+
+def unpatch_gradient_checkpointing():
+    torch.utils.checkpoint = old_checkpoint
+pass
+
+
 # =============================================
 # Fixes Bitsandbytes to remove missing warnings
 from transformers.utils.quantization_config import BitsAndBytesConfig, QuantizationMethod