diff --git a/training.Jenkinsfile b/training.Jenkinsfile
index 8af0cc3..6b78e17 100644
--- a/training.Jenkinsfile
+++ b/training.Jenkinsfile
@@ -1,10 +1,10 @@
 pipeline {
     agent {
-    	dockerfile true
+	    docker {
+	        image 'agakul/ium:mlflow'
+	        args '-v /mlruns:/mlruns'
+	    }
     }
-    	options {
-    		copyArtifactPermission('s444421-evaluation');
-    	}
     	parameters {
     			buildSelector(
 			defaultSelector: lastSuccessful(),
@@ -23,27 +23,12 @@ pipeline {
 				checkout([$class: 'GitSCM', branches: [[name: '*/training_and_evaluation']], extensions: [], userRemoteConfigs: [[credentialsId: 's444421', url: 'https://git.wmi.amu.edu.pl/s444421/ium_444421.git']]])
 			}
 		}
-        stage('Script') {
+        stage('Training') {
             steps {
             	copyArtifacts filter: '*', projectName:'s444421-create-dataset', selector: buildParameter('BUILD_SELECTOR')
-            	sh 'ipython ./training_sacred.py with "epochs=$EPOCHS"'
-            	archiveArtifacts artifacts: 'my_runs/1/*, model.pth, my_runs/_sources/*'
+            	sh 'python training_mlflow.py $EPOCHS'
+            	archiveArtifacts artifacts: 'mlruns/**'
             }
         }
     }
-    post {
-    	success {
-    		emailext body: 'SUCCESS', subject: 's444421-training status', to: 'e19191c5.uam.onmicrosoft.com@emea.teams.ms'
-      		build job: 's444421-evaluation/training_and_evaluation'
-      		}
-      	failure {
-      		emailext body: 'FAILURE', subject: 's444421-training status', to: 'e19191c5.uam.onmicrosoft.com@emea.teams.ms'
-      	}
-      	aborted {
-      		emailext body: 'ABORTED', subject: 's444421-training status', to: 'e19191c5.uam.onmicrosoft.com@emea.teams.ms'
-      	}
-      	changed {
-      		emailext body: 'CHANGED', subject: 's444421-training status', to: 'e19191c5.uam.onmicrosoft.com@emea.teams.ms'
-      	}
-   }
 }
diff --git a/training_mlflow.py b/training_mlflow.py
index 65426f2..4442103 100644
--- a/training_mlflow.py
+++ b/training_mlflow.py
@@ -118,9 +118,9 @@ def my_main(epochs):
     tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
     
     if tracking_url_type_store != "file":
-        mlflow.sklearn.log_model(model, "s444421", registered_model_name="s444421", signature=signature)
+        mlflow.sklearn.log_model(model, "s444421", registered_model_name="s444421", signature=signature, input_example=X_test.numpy()[:5])
     else:
-        mlflow.sklearn.log_model(model, "s444421", signature=signature)
+        mlflow.sklearn.log_model(model, "s444421", signature=signature, input_example=X_test.numpy()[:5])
 
 
 # In[ ]: