diff --git a/JenkinsfileNeural b/JenkinsfileNeural
index e2371e7..232330c 100644
--- a/JenkinsfileNeural
+++ b/JenkinsfileNeural
@@ -16,7 +16,7 @@ node {
 
         }
         stage('Clone repo') {
-                   /* try {*/ docker.image("karopa/ium:30").inside {
+                   /* try {*/ docker.image("karopa/ium:31").inside {
                         stage('Test') {
                             checkout([$class: 'GitSCM', branches: [[name: '*/mlflow']], doGenerateSubmoduleConfigurations: false, extensions: [], submoduleCfg: [], userRemoteConfigs: [[url: 'https://git.wmi.amu.edu.pl/s434765/ium_434765']]])
                             copyArtifacts fingerprintArtifacts: true, projectName: 's434765-create-dataset', selector: buildParameter("BUILD_SELECTOR")
@@ -24,6 +24,7 @@ node {
                                 #!/usr/bin/env bash
                                 chmod 777 neural_network.sh
                                 ./neural_network.sh $EPOCHS | tee output.txt
+                                mlflow run ium_434765 -e main
                             '''
                             archiveArtifacts 'output.txt'
                             archiveArtifacts 'model/**/*.*'
diff --git a/MLproject b/MLproject
index aecf084..af7890e 100644
--- a/MLproject
+++ b/MLproject
@@ -1,7 +1,7 @@
 name: tutorial
 
 docker_env:
-  image: karopa/ium:30
+  image: karopa/ium:31
 
 entry_points:
   main:
diff --git a/neural_network.py b/neural_network.py
index 7de514f..8f93e92 100644
--- a/neural_network.py
+++ b/neural_network.py
@@ -73,4 +73,4 @@ with mlflow.start_run() as run:
                                                                 "channel_title", "views", "likes", "dislikes",
                                                                 "comment_count"]).dropna()
     X_test = data.loc[:, data.columns == "views"].astype(int)
-    mlflow.keras.save_model(model, "model", signature=signature, input_example=X_test)
+    mlflow.keras.save_model(model, "model", registered_model_name="model", signature=signature, input_example=X_test)