diff --git a/vgsales-mlflow.py b/vgsales-mlflow.py
index 18442ff..617d792 100644
--- a/vgsales-mlflow.py
+++ b/vgsales-mlflow.py
@@ -14,10 +14,12 @@ from tensorflow.keras.models import Model
 from tensorflow.keras.callbacks import EarlyStopping
 from tensorflow.keras.models import Sequential
 import mlflow
+from mlflow.tracking import MlflowClient
 
 
 mlflow.set_tracking_uri("http://172.17.0.1:5000")
 mlflow.set_experiment("s434749")
+client = MlflowClient()
 
 def my_main(epochs, batch_size):