fix yolo and text_encoder tests

huggingface · May 3, 2024 · 62f8830 · 62f8830
1 parent 1215075
commit 62f8830
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/tests/models/vision_text_dual_encoder/test_modeling_vision_text_dual_encoder.py b/tests/models/vision_text_dual_encoder/test_modeling_vision_text_dual_encoder.py
@@ -310,6 +310,10 @@ def get_pretrained_model_and_inputs(self):
         return model, inputs
 
     def get_vision_text_model(self, vision_config, text_config):
+        # Run in eager mode if we are in slow mode
+        if _run_slow_tests:
+            vision_config._attn_implementation = "eager"
+
         vision_model = ViTModel(vision_config).eval()
         text_model = BertModel(text_config).eval()
         return vision_model, text_model
@@ -397,6 +401,10 @@ def check_vision_text_output_attention(
         )
 
     def get_vision_text_model(self, vision_config, text_config):
+        # Run in eager mode if we are in slow mode
+        if _run_slow_tests:
+            vision_config._attn_implementation = "eager"
+
         vision_model = DeiTModel(vision_config).eval()
         text_model = RobertaModel(text_config).eval()
         return vision_model, text_model

diff --git a/tests/models/yolos/test_modeling_yolos.py b/tests/models/yolos/test_modeling_yolos.py
@@ -226,6 +226,7 @@ def test_model(self):
     def test_attention_outputs(self):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
         config.return_dict = True
+        config._attn_implementation = "eager"
 
         # in YOLOS, the seq_len is different
         seq_len = self.model_tester.expected_seq_len