diff --git a/Jenkinsfile_train b/Jenkinsfile_train
index dc0985b..8b01bbb 100644
--- a/Jenkinsfile_train
+++ b/Jenkinsfile_train
@@ -1,24 +1,40 @@
 pipeline {
-    agent {dockerfile true}
-	parameters {
-		buildSelector(
-					defaultSelector: lastSuccessful(),
-					description: 'Which build to use for copying artifacts',
-					name: 'BUILD_SELECTOR')
-		string(
-            defaultValue: '500',
-            description: 'Enter the number of Epochs',
-            name: 'epochs',
-            trim: false)
-	}
+    agent {docker { image 'snowycocoon/ium_434788:2' }}
+	parameters{
+        buildSelector(
+            defaultSelector: lastSuccessful(),
+            description: 'Which build to use for copying artifacts',
+            name: 'WHICH_BUILD'
+        )
+        string(
+            defaultValue: '16',
+            description: 'batch size',
+            name: 'BATCH_SIZE'
+        )
+        string(
+            defaultValue: '15', 
+            description: 'epochs',
+            name: 'EPOCHS'
+        )
+    }
 
     stages {
-		stage('copyArtifacts') {
-            steps {
-                copyArtifacts fingerprintArtifacts: true, projectName: 's434788-create-dataset', selector: buildParameter('BUILD_SELECTOR')
-				sh 'python3 Zadanie_06_training.py ${epochs}'
+		stage('copy artifacts')
+            {
+                steps
+                {
+                    copyArtifacts(fingerprintArtifacts: true, projectName: 's434788-create-dataset', selector: buildParameter('WHICH_BUILD'))
+                }
+            }
+            stage('train')
+            {
+                steps
+                {
+                    catchError {
+                        sh 'python3.8 Zadanie_06_training'
+                    }
+                }
             }
-        }
 
 		stage('Archive artifacts') {
 			steps{