diff --git a/generate_MLmodel.py b/generate_MLmodel.py
index 8678258..77c6098 100644
--- a/generate_MLmodel.py
+++ b/generate_MLmodel.py
@@ -29,6 +29,7 @@ if __name__ == "__main__":
 
     lr = args.lr
     n_epochs = args.epochs
+    mlflow.set_tracking_uri("http://127.0.0.1:5000")
     mlflow.set_experiment("s426206")
     with mlflow.start_run():
         mlflow.log_param("lr", lr)
@@ -129,7 +130,7 @@ if __name__ == "__main__":
             siganture = infer_signature(x_train, model(torch.tensor(np.reshape(x_train, (-1,1))).float()).numpy())
 
         #mlflow.set_experiment("s426206")
-        mlflow.set_tracking_uri("http://172.17.0.1:5000")
+        #mlflow.set_tracking_uri("http://172.17.0.1:5000")
         tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
         # print(tracking_url_type_store)
         # Model registry does not work with file store