Changing batch size

actualize-ae · Dec 26, 2023 · 21343f4 · 21343f4
1 parent d2837a0
commit 21343f4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/start.py b/start.py
@@ -67,7 +67,7 @@ def run_vllm(model_id_or_path, tasks, shots):
     cmd = f"lm_eval --model=vllm \
                     --model_args={model_args_str} \
                     --tasks={tasks} \
-                    --batch_size=8 \
+                    --batch_size=4 \
                     --num_fewshot={shots} \
                     --output_path=/opt/ml/model/results_{shots}_shots.json"
     print(f"Running command: {cmd}")
@@ -85,7 +85,7 @@ def run_hf(model_id_path, peft_model_id_or_path, tasks, shots):
     cmd = f"lm_eval --model hf \
                     --model_args {model_args_str} \
                     --tasks {tasks} \
-                    --batch_size=8 \
+                    --batch_size=4 \
                     --num_fewshot={shots} \
                     --output_path=/opt/ml/model/results_{shots}_shots.json"
     print(f"Running command: {cmd}")