File tree Expand file tree Collapse file tree 1 file changed +5
-2
lines changed
src/transformers/generation/continuous_batching Expand file tree Collapse file tree 1 file changed +5
-2
lines changed Original file line number Diff line number Diff line change 2222from itertools import count
2323from math import ceil
2424from time import perf_counter
25- from tqdm .contrib .logging import logging_redirect_tqdm
2625
2726import torch
2827from torch import nn
2928from tqdm import tqdm
29+ from tqdm .contrib .logging import logging_redirect_tqdm
3030
3131from ...configuration_utils import PretrainedConfig
3232from ...generation .configuration_utils import GenerationConfig
@@ -810,6 +810,7 @@ def is_running(self) -> bool:
810810 """Check if the background generation thread is running."""
811811 return self ._generation_thread is not None and self ._generation_thread .is_alive ()
812812
813+ # NOTE: don't forget to update `continuous_batching_context_manager` when changing this method's definition
813814 def stop (self , block : bool = True , timeout : float | None = None ) -> None :
814815 """Signal the background thread to stop.
815816
@@ -1087,7 +1088,9 @@ def continuous_batching_context_manager(
10871088 try :
10881089 yield manager
10891090 finally :
1090- logger .debug ("Continuous batching loop finished" ) # a dummy log needed for the logs of stop to show. Won't show
1091+ logger .debug (
1092+ "Continuous batching loop finished"
1093+ ) # a dummy log needed for the logs of stop to show. Won't show
10911094 manager .stop (block = block , timeout = timeout )
10921095
10931096 # NOTE: don't forget to update `continuous_batching_context_manager` when changing this method's definition
You can’t perform that action at this time.
0 commit comments