diff --git a/pytorch/pytorch.py b/pytorch/pytorch.py
index 5dbd84b..e54a5f1 100644
--- a/pytorch/pytorch.py
+++ b/pytorch/pytorch.py
@@ -24,6 +24,7 @@ from sacred.observers import MongoObserver
 import mlflow
 import mlflow.pytorch
 from mlflow.models import infer_signature
+from urllib.parse import urlparse
 
 # In[2]:
 ex = Experiment(save_git_info=False)