diff --git a/Dockerfile b/Dockerfile
index 39f8147..437e4b4 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -6,14 +6,15 @@ RUN python3 -m pip install kaggle
 RUN python3 -m pip install pandas
 RUN python3 -m pip install scikit-learn
 RUN python3 -m pip install --upgrade tensorflow
+RUN python3 -m pip install dvc
+RUN python3 -m pip install dvc[ssh] paramiko
 
 WORKDIR app
 
 ARG EPOCHS
 ENV EPOCHS=${EPOCHS}
 
-COPY mlflow_training.py ./
-COPY heart_2020_cleaned.csv ./
+COPY ml_prepare.py ./
+COPY ml_training.py ./
 
-
-CMD ["python3", "./mlflow_training.py $EPOCHS"]
+CMD ["python3", "./ml_training.py $EPOCHS"]
diff --git a/Jenkinsfile.dvc_repo b/Jenkinsfile.dvc_repo
index b3756e1..401a750 100644
--- a/Jenkinsfile.dvc_repo
+++ b/Jenkinsfile.dvc_repo
@@ -22,7 +22,7 @@ pipeline {
                     sh 'dvc remote add -d ium_ssh_remote ssh://ium-sftp@tzietkiewicz.vm.wmi.amu.edu.pl/ium-sftp'
                     sh 'dvc remote modify --local ium_ssh_remote keyfile $IUM_SFTP_KEY'
                     sh 'dvc pull'
-                    }
+                }
             }
         }
     }