Update experiment.py

allenai · Jul 24, 2024 · d5f7b85 · d5f7b85
1 parent c906ae7
commit d5f7b85
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/make_it_batch/experiment.py b/tests/make_it_batch/experiment.py
@@ -274,7 +274,7 @@ def tag(self):
  def training_pipeline(self, **kwargs):
  ppo_steps = int(1e6)
  lr = 2.5e-4
- num_mini_batch = 2 if not torch.cuda.is_available() else 6
+ num_mini_batch = self.task_batch_size if not torch.cuda.is_available() else 6
  update_repeats = 4
  num_steps = self.MAX_STEPS
  metric_accumulate_interval = self.MAX_STEPS * 1 # Log every 10 max length tasks