From 76b2474be865e0c5b716ba169f1822b90dd73304 Mon Sep 17 00:00:00 2001 From: Dorian Hoxha Date: Mon, 25 Mar 2024 12:52:16 +0100 Subject: [PATCH] Make tqdm optional cause it interferes with thread monkey patching in gevent --- dl_translate/_translation_model.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/dl_translate/_translation_model.py b/dl_translate/_translation_model.py index 1e65b16..7d532b0 100644 --- 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)