Correct @is_flaky test decoration (huggingface#31480)

* Correct @is_flaky decorator
robinderat · Jun 24, 2024 · 3c2d4d6 · 3c2d4d6
1 parent 4b82256
commit 3c2d4d6
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/tests/models/gemma/test_modeling_gemma.py b/tests/models/gemma/test_modeling_gemma.py
@@ -493,7 +493,7 @@ def test_sdpa_equivalence(self):
     @require_flash_attn
     @require_torch_gpu
     @pytest.mark.flash_attn_test
-    @is_flaky
+    @is_flaky()
     @slow
     def test_flash_attn_2_equivalence(self):
         for model_class in self.all_model_classes:

diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
@@ -3407,7 +3407,7 @@ def test_flash_attn_2_conversion(self):
     @require_torch_gpu
     @mark.flash_attn_test
     @slow
-    @is_flaky
+    @is_flaky()
     def test_flash_attn_2_inference_equivalence(self):
         for model_class in self.all_model_classes:
             if not model_class._supports_flash_attn_2:
@@ -3501,7 +3501,7 @@ def test_flash_attn_2_inference_equivalence(self):
     @require_torch_gpu
     @mark.flash_attn_test
     @slow
-    @is_flaky
+    @is_flaky()
     def test_flash_attn_2_inference_equivalence_right_padding(self):
         for model_class in self.all_model_classes:
             if not model_class._supports_flash_attn_2:
@@ -3591,7 +3591,7 @@ def test_flash_attn_2_inference_equivalence_right_padding(self):
     @require_torch_gpu
     @mark.flash_attn_test
     @slow
-    @is_flaky
+    @is_flaky()
     def test_flash_attn_2_generate_left_padding(self):
         for model_class in self.all_generative_model_classes:
             if not model_class._supports_flash_attn_2:
@@ -3635,7 +3635,7 @@ def test_flash_attn_2_generate_left_padding(self):
     @require_flash_attn
     @require_torch_gpu
     @mark.flash_attn_test
-    @is_flaky
+    @is_flaky()
     @slow
     def test_flash_attn_2_generate_padding_right(self):
         for model_class in self.all_generative_model_classes: