fix test case

Former-commit-id: b075b2971c6acb2c6039b36420a296f1f4e1b91b
2024-09-08 01:50:51 +08:00
parent 158e0e1f63
commit b6810b209a
2 changed files with 4 additions and 1 deletions
--- a/tests/model/model_utils/test_checkpointing.py
+++ b/tests/model/model_utils/test_checkpointing.py
@@ -54,7 +54,7 @@ def test_checkpointing_disable():
 def test_unsloth_gradient_checkpointing():
    model = load_train_model(use_unsloth_gc=True, **TRAIN_ARGS)
    for module in filter(lambda m: hasattr(m, "gradient_checkpointing"), model.modules()):
-        assert module._gradient_checkpointing_func.__self__.__name__ == "UnslothGradientCheckpointing"
+        assert module._gradient_checkpointing_func.__self__.__name__ == "UnslothGradientCheckpointing"  # classmethod


 def test_upcast_layernorm():