fix test case

Former-commit-id: b332908ab4
2025-12-16 03:40:34 +08:00 · 2024-09-08 01:50:51 +08:00
parent 945841503e
commit f6f58ebef0
2 changed files with 4 additions and 1 deletions
--- a/src/llamafactory/model/model_utils/checkpointing.py
+++ b/src/llamafactory/model/model_utils/checkpointing.py
@@ -89,6 +89,9 @@ def get_custom_gradient_checkpointing_func(gradient_checkpointing_func):

        return gradient_checkpointing_func(func, *args, **kwargs)

+    if hasattr(gradient_checkpointing_func, "__self__"):  # fix test case
+        custom_gradient_checkpointing_func.__self__ = gradient_checkpointing_func.__self__
+
    return custom_gradient_checkpointing_func