From 838a0268b88b6e56783b401f9cacae9cb0cbb120 Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Tue, 8 Jul 2025 15:36:05 +0200 Subject: [PATCH] fix flaky `test_generate_compile_model_forward` (#39276) fix Co-authored-by: ydshieh --- tests/generation/test_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index 9edb1fe99f..e01b1ac50b 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -2088,6 +2088,7 @@ class GenerationTesterMixin: ⚠️ Runs two sequential generations to ensure the cache doesn't get stuck after the first compiled run! ⚠️ """ + set_model_tester_for_less_flaky_test(self) for model_class in self.all_generative_model_classes: # 1. Test exclusion criteria if not model_class._supports_static_cache: @@ -2095,7 +2096,9 @@ class GenerationTesterMixin: # 2. Prepares two sets of inputs config, inputs_dict = self.prepare_config_and_inputs_for_generate(batch_size=4) + set_config_for_less_flaky_test(config) model = model_class(config).to(torch_device) + set_model_for_less_flaky_test(model) model.eval() # otherwise `self.training` is `True` -- this flag is used at attn mask creation time # Some composite models have a custom generate and will call an inner model's generate -> that inner model