From 39c4dc62e7cdfd4d2e63eb6955621474a467ccb0 Mon Sep 17 00:00:00 2001 From: Nadav Elyahu Date: Wed, 4 Sep 2024 12:25:05 +0300 Subject: [PATCH] fix pipeline eval_batch micro_batches argument for schedule --- deepspeed/runtime/pipe/engine.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py index 26196ff37ac4..7ebf5487cf9e 100644 --- a/deepspeed/runtime/pipe/engine.py +++ b/deepspeed/runtime/pipe/engine.py @@ -482,9 +482,7 @@ def eval_batch(self, micro_batches = self.micro_batches if num_micro_batches is None else num_micro_batches # Do the work - sched = schedule.InferenceSchedule(micro_batches=self.micro_batches, - stages=self.num_stages, - stage_id=self.stage_id) + sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=self.num_stages, stage_id=self.stage_id) # prevent dead-lock with multiple evals sequence dist.barrier()