fix gptq tests

huggingface · Nov 23, 2024 · 5e5aa3f · 5e5aa3f
1 parent 1945f91
commit 5e5aa3f
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py
@@ -1196,7 +1196,17 @@ def require_auto_gptq(test_case):
     """
     Decorator for auto_gptq dependency
     """
-    return unittest.skipUnless(is_auto_gptq_available(), "test requires auto-gptq")(test_case)
+    eetq_available = is_eetq_available()
+    if eetq_available:
+        try:
+            import eetq  # noqa: F401
+        except ImportError as exc:
+            if "shard_checkpoint" in str(exc):
+                # EETQ 1.0.0 is currently broken with the latest transformers because it tries to import the removed
+                # shard_checkpoint function, see https://github.com/NetEase-FuXi/EETQ/issues/34.
+                # TODO: Remove once eetq releases a fix and this release is used in CI
+                eetq_available = False
+    return unittest.skipUnless(is_auto_gptq_available() and eetq_available, "test requires auto-gptq")(test_case)
 
 
 def require_auto_awq(test_case):