diff --git a/Jenkinsfile-train b/Jenkinsfile-train
index eb05d97..1966e06 100644
--- a/Jenkinsfile-train
+++ b/Jenkinsfile-train
@@ -36,9 +36,6 @@ 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 2983507..edf5521 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")
 
 
@@ -80,8 +80,8 @@ if __name__ == "__main__":
                 registered_model_name="s444409",
                 signature=signature,
                 input_example=input_example,
-                extra_files=['model.py']
+                code_paths=['model.py']
             )
         else:
             mlflow.pytorch.log_model(model, "s444409", signature=signature, input_example=input_example,
-                                     extra_files=['model.py'])
+                                     code_paths=['model.py'])