From 4b30c7a89b9bcf5e5ddb64fcfeddac19ad73d24c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 17 Nov 2024 17:12:34 -0800 Subject: [PATCH] Update _utils.py --- unsloth/models/_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 629cee1c7e..7243d34c3a 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -57,6 +57,7 @@ __all__ = [ "fused_linear_cross_entropy", "patch_unsloth_smart_gradient_checkpointing", "unpatch_unsloth_smart_gradient_checkpointing", + "create_gradient_checkpointing_buffer", "patch_compiled_autograd", "process_vision_info", @@ -93,6 +94,7 @@ from unsloth_zoo.gradient_checkpointing import ( patch_unsloth_smart_gradient_checkpointing, unpatch_unsloth_smart_gradient_checkpointing, + create_gradient_checkpointing_buffer, ) from unsloth_zoo.loss_utils import ( HAS_CUT_CROSS_ENTROPY,