From e5d448b2f6f39d6ac8b0f40d828eaafb3cc009df Mon Sep 17 00:00:00 2001
From: Michal Zukowski <michal.zukowski@reef.pl>
Date: Thu, 26 Sep 2024 13:01:36 +0200
Subject: [PATCH] Fix handling leftover prompts

---
 src/compute_horde_prompt_gen/run.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/compute_horde_prompt_gen/run.py b/src/compute_horde_prompt_gen/run.py
index b2ccfd3..0034b0b 100644
--- a/src/compute_horde_prompt_gen/run.py
+++ b/src/compute_horde_prompt_gen/run.py
@@ -57,9 +57,8 @@ def generate_prompts(
         )
 
         # Use leftover prompts from previous batch if available
-        while leftover_prompts and total_prompts > 0:
+        while leftover_prompts and len(new_prompts) < total_prompts:
             new_prompts.append(leftover_prompts.popleft())
-            total_prompts -= 1
 
         if len(new_prompts) > total_prompts:
             # Save extra prompts for next batch