diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index 2f8e60c79151e9..803942cb8c9841 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -62,7 +62,7 @@ SpeechEncoderDecoderModel, T5ForConditionalGeneration, ) - from transformers.cache_utils import DynamicCache, EncoderDecoderCache, QuantoQuantizedCache, StaticCache + from transformers.cache_utils import DynamicCache, DynamicSlidingWindowCache, EncoderDecoderCache, QuantoQuantizedCache, StaticCache from transformers.generation import ( BeamSampleDecoderOnlyOutput, BeamSampleEncoderDecoderOutput, @@ -1868,6 +1868,9 @@ def test_new_cache_format(self, num_beams, do_sample): if config.is_encoder_decoder: cache_cls = EncoderDecoderCache past_key_values = cache_cls(DynamicCache(), DynamicCache()) + elif getattr(self.config, "sliding_window", None) is not None: + cache_cls = DynamicSlidingWindowCache + past_key_values = cache_cls(self.config.sliding_window) else: cache_cls = DynamicCache past_key_values = cache_cls()