From 4f2e6b267af34a7036476f029b5ac8c62c256a44 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Fri, 20 Sep 2024 18:25:25 +0200 Subject: [PATCH] update cache conversion in test --- tests/generation/test_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index 6370457917fcaf..f3bffd612a602e 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -1907,7 +1907,10 @@ def test_new_cache_format(self, num_beams, do_sample): ) new_cache = new_results.past_key_values - legacy_cache_converted = cache_cls.from_legacy_cache(legacy_results.past_key_values) + if cache_cls == DynamicSlidingWindowCache: + legacy_cache_converted = cache_cls.from_legacy_cache(config.sliding_window, legacy_results.past_key_values) + else: + legacy_cache_converted = cache_cls.from_legacy_cache(legacy_results.past_key_values) for layer_idx in range(len(new_cache)): for kv_idx in range(len(new_cache[layer_idx])): self.assertTrue(