fix tqdm on colab

hse-aml · Dec 11, 2018 · 2bd4ff5 · 2bd4ff5
1 parent cb87534
commit 2bd4ff5
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/tqdm_utils.py b/tqdm_utils.py
@@ -47,9 +47,17 @@ def __next__(self):
             raise StopIteration
 
 
-def tqdm_notebook_failsafe(*args, **kwargs):
+def is_colab():
     try:
-        return tqdm.tqdm_notebook(*args, **kwargs)
-    except:
+        import google.colab
+        return True
+    except ImportError:
+        return False
+
+
+def tqdm_notebook_failsafe(*args, **kwargs):
+    if is_colab():
         # tqdm is broken on Google Colab
         return SimpleTqdm(*args, **kwargs)
+    else:
+        return tqdm.tqdm_notebook(*args, **kwargs)