diff --git a/Jenkinsfile-train b/Jenkinsfile-train
index 1966e06..eb05d97 100644
--- a/Jenkinsfile-train
+++ b/Jenkinsfile-train
@@ -36,6 +36,9 @@ pipeline {
         stage('Train model') {
             steps {
                 sh "python train_model.py -e ${params.EPOCHS} -b ${params.BATCHSIZE}"
+                archiveArtifacts artifacts: 'model_out', onlyIfSuccessful: true
+                archiveArtifacts artifacts: 'mlruns/**', onlyIfSuccessful: true
+                sh 'rm -r mlruns'
             }
         }
         stage('Evaluate model') {
diff --git a/train_model.py b/train_model.py
index edf5521..8a8aa06 100644
--- a/train_model.py
+++ b/train_model.py
@@ -14,7 +14,7 @@ default_epochs = 5
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-mlflow.set_tracking_uri("http://172.17.0.1:5000")
+# mlflow.set_tracking_uri("http://172.17.0.1:5000")
 mlflow.set_experiment("s444409")