[CI] green llama tests (#37244)

* green llama tests * use cleanup instead * better test comment; cleanup upgrade * better test comment; cleanup upgrade
2025-04-03 14:15:53 +01:00
parent 782d7d945d
commit 9a1c1fe7ed
15 changed files with 62 additions and 36 deletions
--- a/tests/quantization/aqlm_integration/test_aqlm.py
+++ b/tests/quantization/aqlm_integration/test_aqlm.py
@@ -226,7 +226,7 @@ class AqlmTest(unittest.TestCase):
        # Setup static KV cache for generation
        past_key_values = StaticCache(
            config=self.quantized_model.config,
-            batch_size=1,
+            max_batch_size=1,
            max_cache_len=seq_length + self.max_new_tokens + 1,
            device=torch_device,
            dtype=self.quantized_model.config._pre_quantization_dtype,