[#4128] temp fix

bentoml · Aug 20, 2023 · 6743e12 · 6743e12
1 parent 47241b1
commit 6743e12
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/src/bentoml/_internal/marshal/dispatcher.py b/src/bentoml/_internal/marshal/dispatcher.py
@@ -193,12 +193,6 @@ async def train_optimizer(
         if self.max_batch_size < batch_size:
             batch_size = self.max_batch_size
 
-        if batch_size > 1:
-            wait = min(
-                self.max_latency * 0.95,
-                (batch_size * 2 + 1) * (self.optimizer.o_a + self.optimizer.o_b),
-            )
-
         req_count = 0
         try:
             while req_count < num_reqs_to_train:
@@ -217,6 +211,10 @@ async def train_optimizer(
                 if batch_size > 1:  # only wait if batch_size
                     a = self.optimizer.o_a
                     b = self.optimizer.o_b
+                    wait = min(
+                        self.max_latency * 0.95,
+                        (batch_size * 2 + 1) * (self.a + self.optimizer.o_b),
+                    )
 
                     if n < batch_size and (batch_size * a + b) + w0 <= wait:
                         await asyncio.sleep(self.tick_interval)