Set the batch_size

examples/benchmark.py

@@ -43,7 +43,7 @@ def run_round(model_path, quant_file, n_generate, input_ids, batch_size):
     print(f" -- Loading model...")
     model = AutoAWQForCausalLM.from_quantized(
         model_path, quant_file, fuse_layers=True,
-        max_new_tokens=n_generate
+        max_new_tokens=n_generate, batch_size=batch_size
     )
 
     print(f" -- Warming up...")