Mark test_constrained_beam_search_generate as flaky (huggingface#28757)

* Make test_constrained_beam_search_generate as flaky * Update tests/generation/test_utils.py
SangbumChoi · Jan 31, 2024 · 2e865a4 · 2e865a4
1 parent 7e6a905
commit 2e865a4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py
@@ -1293,6 +1293,8 @@ def test_group_beam_search_generate_dict_output(self):
                     output, input_ids, model.config, num_return_sequences=num_return_sequences * beam_scorer.num_beams
                 )
 
+    # TODO: @gante
+    @is_flaky()
     def test_constrained_beam_search_generate(self):
         for model_class in self.all_generative_model_classes:
             config, input_ids, attention_mask, max_length = self._get_input_ids_and_config()