From 2e865a4704388e621bfafe3c0b154ef5af87c45d Mon Sep 17 00:00:00 2001 From: amyeroberts <22614925+amyeroberts@users.noreply.github.com> Date: Mon, 29 Jan 2024 15:22:25 +0000 Subject: [PATCH] Mark test_constrained_beam_search_generate as flaky (#28757) * Make test_constrained_beam_search_generate as flaky * Update tests/generation/test_utils.py --- tests/generation/test_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index 2c16f41ae171dc..855187778d2cf0 100644 --- 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()