added batch size to SUT Server (#1750)

mlcommons · Jun 25, 2024 · a31cced · a31cced
1 parent eb6714e
commit a31cced
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/language/llama2-70b/SUT.py b/language/llama2-70b/SUT.py
@@ -285,7 +285,7 @@ def __del__(self):
 
 
 class SUTServer(SUT):
-    def __init__(self, model_path=None, dtype="bfloat16", device="cpu", total_sample_count=24576, dataset_path=None, workers=1):
+    def __init__(self, model_path=None, dtype="bfloat16", device="cpu", total_sample_count=24576, dataset_path=None, batch_size=None, workers=1):
 
         super().__init__(model_path=model_path, dtype=dtype, device=device, total_sample_count=total_sample_count, dataset_path=dataset_path, workers=workers)