diff --git a/Jenkins_train b/Jenkins_train
index 84248ce..6170292 100644
--- a/Jenkins_train
+++ b/Jenkins_train
@@ -1,8 +1,5 @@
 pipeline {
     agent any
-    environment {
-        dockerImage = ''
-    }
     parameters {
         string(name: 'EPOCHS', defaultValue: '100', description: 'Number of epochs')
         string(name: 'LR', defaultValue: '0.01', description: 'Learning rate')
@@ -13,28 +10,16 @@ pipeline {
                 checkout([$class: 'GitSCM', branches: [[name: '*/master']], extensions: [], userRemoteConfigs: [[credentialsId: 's487194', url: 'https://git.wmi.amu.edu.pl/s487194/ium_487194']]])
             }
         }
-        stage('Pull Docker image') {
+        stage('Docker and Training'){
             steps {
                 script {
-                    dockerImage = docker.image('s487194/ium:lab5')
-                }
-            }
-        }
-        stage('Train model') {
-            steps {
-                script {
-                    dockerImage.inside('-u root') {
+                    def dockerImage = docker.image('s487194/ium:lab5')
+                    dockerImage.inside { 
                         sh "python train.py -epochs ${params.EPOCHS} -lr ${params.LR}"
+                        archiveArtifacts artifacts: 'classificationn_model.pt', onlyIfSuccessful: true
                     }
                 }
             }
         }
-        stage('Archive model') {
-            steps {
-                archiveArtifacts artifacts: 'classificationn_model.pt', onlyIfSuccessful: true
-            }
-        }
     }
 }
-
-