diff --git a/Dockerfile b/Dockerfile
index 35cd423..ab6509a 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -19,4 +19,8 @@ RUN pip3 install --user keras
 RUN pip3 install --user GitPython
 RUN pip3 install --user pymongo
 
+WORKDIR /app
 
+COPY ./train.py ./
+
+CMD ./train.py
diff --git a/Jenkinsfile b/Jenkinsfile
index 9af792b..e8105bd 100644
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -1,5 +1,5 @@
 pipeline {
-    agent any;
+    agent { dockerfile true }
 
    parameters{
       buildSelector(
@@ -31,9 +31,7 @@ pipeline {
             {
                 steps
                 {
-			withDockerContainer('shroomy/ium2:2') {
-                       	sh 'python3 train.py ${BATCH_SIZE} ${EPOCHS}'
-			}
+                       sh 'python3 train.py ${BATCH_SIZE} ${EPOCHS}'
                 }
             }