diff --git a/Jenkinsfile_artifact b/Jenkinsfile_artifact
index 52a317e..0f33409 100644
--- a/Jenkinsfile_artifact
+++ b/Jenkinsfile_artifact
@@ -1,9 +1,9 @@
 pipeline {
     agent {
-        dockerfile {
-            args '-v /mlruns:/mlruns'
+        docker {
+            image 'ksero/ium:pytorch'
         }
-    }
+    } 
     parameters {
         buildSelector(
 			defaultSelector: lastSuccessful(),
diff --git a/Jenkinsfile_mlflow b/Jenkinsfile_mlflow
index 478b5ec..9f854a4 100644
--- a/Jenkinsfile_mlflow
+++ b/Jenkinsfile_mlflow
@@ -1,9 +1,7 @@
 pipeline {
-    agent {
-        dockerfile {
-            args '-v /mlruns:/mlruns'
+    docker {
+            image 'ksero/ium:pytorch'
         }
-    }    
     stages {
         stage('Prediction') {
             steps {
diff --git a/MLproject b/MLproject
index 66cfaa1..86a8c57 100644
--- a/MLproject
+++ b/MLproject
@@ -2,7 +2,7 @@ name: tutorial
 # conda_env: conda.yaml 
     
 docker_env:
-  image: ksero/ium:mlflow
+  image: ksero/ium:pytorch
 
 entry_points:
   main: