Remove deprecated batch_size parameter (#37007)

2025-03-27 23:01:56 +08:00
parent 4cc65e990f
commit 6cc9c8d7d1
4 changed files with 50 additions and 97 deletions
--- a/tests/models/mamba/test_modeling_mamba.py
+++ b/tests/models/mamba/test_modeling_mamba.py
@@ -422,7 +422,7 @@ class MambaModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixi
        model.eval()

        # Create cache with float32 dtype
-        cache_params = MambaCache(config, batch_size=input_ids.size(0), dtype=torch.float32, device=torch_device)
+        cache_params = MambaCache(config, max_batch_size=input_ids.size(0), dtype=torch.float32, device=torch_device)

        # If code is correct, no error occurs and test passes
        outputs = model(