add batch size

mikecovlee · Jul 18, 2024 · 4a98493 · 4a98493
1 parent 41b4410
commit 4a98493
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/mlora/generator.py b/mlora/generator.py
@@ -208,6 +208,8 @@ def _dispatch_task_in(
         else:
             per_task_jobs = concurrent_jobs
 
+        per_task_jobs = min(per_task_jobs, config.batch_size)
+
         batch_start_idx = len(input_tokens)
         while per_task_jobs > 0 and len(config.data_) > 0:
             per_task_jobs = per_task_jobs - 1