Modify num_steps to num_steps_remaining

kcmclau21122 · Jun 9, 2024 · c3bd6bc · c3bd6bc
1 parent fc28ef9
commit c3bd6bc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimization_pipeline.py b/optimization_pipeline.py
@@ -270,7 +270,7 @@ def run_pipeline(self, num_steps: int):
         # Run the optimization pipeline for num_steps
         num_steps_remaining = num_steps - self.batch_id
         for i in range(num_steps_remaining):
-            stop_criteria = self.step(i, num_steps)
+            stop_criteria = self.step(i, num_steps_remaining)
             if stop_criteria:
                 break
         final_result = self.extract_best_prompt()