add comment back

huggingface · Feb 7, 2024 · 2c650c9 · 2c650c9
1 parent b13b302
commit 2c650c9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/lighteval/models/base_model.py b/src/lighteval/models/base_model.py
@@ -364,6 +364,8 @@ def greedy_until(
                 max_input_length=max_context_continuation_size_allowed,
                 starting_batch_size=starting_batch_size,
             )
+            # For next iteration, since the batch will be smaller, we'll test a bigger batch size
+            starting_batch_size = batch_size * 2
 
             dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=lambda batch: batch)
             if self.accelerator: