Make tqdm optional cause it interferes with thread monkey patching in…

… gevent
xhluca · Mar 25, 2024 · 76b2474 · 76b2474
1 parent 4ed0b8d
commit 76b2474
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/dl_translate/_translation_model.py b/dl_translate/_translation_model.py
@@ -177,8 +177,11 @@ def translate(
         data_loader = torch.utils.data.DataLoader(text, batch_size=batch_size)
         output_text = []
 
+        tqdm_iterator = data_loader
+        if verbose is True:
+            tqdm_iterator = tqdm
         with torch.no_grad():
-            for batch in tqdm(data_loader, disable=not verbose):
+            for batch in tqdm_iterator(data_loader, disable=not verbose):
                 encoded = self._tokenizer(batch, return_tensors="pt", padding=True)
                 encoded.to(self.device)