fix flaky test_generate_compile_model_forward (#39276)

fix

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar
2025-07-08 15:36:05 +02:00
committed by GitHub
parent 29d0030e23
commit 838a0268b8

View File

@@ -2088,6 +2088,7 @@ class GenerationTesterMixin:
⚠️ Runs two sequential generations to ensure the cache doesn't get stuck after the first compiled run! ⚠️ ⚠️ 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: for model_class in self.all_generative_model_classes:
# 1. Test exclusion criteria # 1. Test exclusion criteria
if not model_class._supports_static_cache: if not model_class._supports_static_cache:
@@ -2095,7 +2096,9 @@ class GenerationTesterMixin:
# 2. Prepares two sets of inputs # 2. Prepares two sets of inputs
config, inputs_dict = self.prepare_config_and_inputs_for_generate(batch_size=4) 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) 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 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 # Some composite models have a custom generate and will call an inner model's generate -> that inner model