diff --git a/JenkinsfileNeural b/JenkinsfileNeural
index 75d0496..dd93fd9 100644
--- a/JenkinsfileNeural
+++ b/JenkinsfileNeural
@@ -16,7 +16,7 @@ node {
 
         }
         stage('Clone repo') {
-                   /* try {*/ docker.image("karopa/ium:27").inside {
+                   /* try {*/ docker.image("karopa/ium:28").inside {
                         stage('Test') {
                             checkout([$class: 'GitSCM', branches: [[name: '*/master']], 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")
diff --git a/neural_network.py b/neural_network.py
index fbf75c1..5e437c1 100644
--- a/neural_network.py
+++ b/neural_network.py
@@ -5,8 +5,8 @@ import numpy as np
 
 from tensorflow import keras
 import sys
-import mlflow.sklearn
-
+import mlflow
+import mlflow.models
 import logging
 
 from evaluate_network import evaluate_model
@@ -66,3 +66,11 @@ with mlflow.start_run() as run:
 
     error = evaluate_model()
     mlflow.log_metric("rmse", error)
+    signature = mlflow.models.signature.infer_signature(X, model.predict(y))
+    data = pd.read_csv("data_dev", sep=',', error_bad_lines=False,
+                       skip_blank_lines=True, nrows=527, names=["video_id", "last_trending_date",
+                                                                "publish_date", "publish_hour", "category_id",
+                                                                "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)