diff --git a/cedr/train.py b/cedr/train.py
index e0363cf..ff3c63c 100644
--- a/cedr/train.py
+++ b/cedr/train.py
@@ -3,7 +3,7 @@
 import subprocess
 import random
 import tempfile
-from tqdm import tqdm
+from tqdm.autonotebook import tqdm
 import torch
 from . import modeling
 from . import data