From 1f0fd97f46a740ca5ae42b2bb128273a46e3037d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 4 Nov 2024 01:20:31 -0800 Subject: [PATCH] Update _utils.py --- unsloth/models/_utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 5c099548d6..7a8dcedada 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -41,14 +41,17 @@ __all__ = [ "torch_amp_custom_bwd", "accelerate_old_send_to_device", "accelerate_new_send_to_device", - "patch_gradient_checkpointing", - "unpatch_gradient_checkpointing", "patch_gradient_accumulation_fix", "patch_compiling_bitsandbytes", "patch_regional_compilation", "patch_layernorm", "patch_torch_compile", "patch_model_and_tokenizer", + + "patch_unsloth_gradient_checkpointing", + "unpatch_unsloth_gradient_checkpointing", + "patch_gradient_checkpointing", + "unpatch_gradient_checkpointing", ] import torch