diff --git a/Jenkinsfile_mlflow b/Jenkinsfile_mlflow
index fb51521..4c9bd29 100644
--- a/Jenkinsfile_mlflow
+++ b/Jenkinsfile_mlflow
@@ -1,6 +1,7 @@
 pipeline {
     agent { 
         dockerfile  true
+        args '-v /mlruns:/mlruns'
     }
 
     parameters {
@@ -30,7 +31,6 @@ pipeline {
                 "KAGGLE_KEY=${params.KAGGLE_KEY}"]) {
                     sh 'python3 ./nn_train_mlflow.py'
                     archiveArtifacts artifacts: 'mlruns/**'
-                    sh 'rm -r mlruns'
                     sh 'rm -r my_model_mlflow'
                 }
             }
diff --git a/nn_train_mlflow.py b/nn_train_mlflow.py
index 33839bc..653b0ac 100644
--- a/nn_train_mlflow.py
+++ b/nn_train_mlflow.py
@@ -9,8 +9,11 @@ from keras.utils import np_utils
 from tensorflow import keras
 import mlflow
 import sys
+from urllib.parse import urlparse
+
 
 mlflow.set_experiment("s444517")
+mlflow.set_tracking_uri("http://172.17.0.1:5000")
 
 # reading data
 def read_data():
@@ -59,7 +62,7 @@ with mlflow.start_run():
     second_activation_funct = int(sys.argv[3]) if len(sys.argv) > 3 else "softmax"
 
     x_train_set, dummy_y, x_validate_set, dummy_yv, x_test_set,  y_test_set, y_class_names = data_prep()
-
+    
     number_of_classes = 33
     number_of_features = 5
     model = Sequential()
@@ -79,10 +82,27 @@ with mlflow.start_run():
         y_true.append(sorted(y_class_names)[np.argmax(single_pred)])
         y_pred.append(y_test_set[numerator])
 
+
+    signature = mlflow.models.signature.infer_signature(x_train_set, model.predict(x_train_set))
+    
+    input_example = {
+        "Rating": 4.100000,
+        "Reviews": 0.000001,
+        "Installs": 0.000005,
+        "Price": 0.000000,
+        "Genres_numeric_value": 57.000000
+
+    }
+
     mlflow.log_param("epoch", epoch)
     mlflow.log_param("1st_activation_funct", first_activation_funct)
     mlflow.log_param("2nd_activation_funct", second_activation_funct)
-    #mlflow.keras.log_model(model, 'my_model')
-    mlflow.keras.save_model(model, "my_model_mlflow")
     mlflow.log_metric("accuracy", accuracy_score(y_true, y_pred))
-    
\ No newline at end of file
+
+    tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
+    
+    if tracking_url_type_store != "file":
+        mlflow.sklearn.log_model(model, "my_model_mlflow", registered_model_name="s444517", signature=signature, input_example=input_example)
+    else:
+        mlflow.sklearn.log_model(model, "my_model_mlflow", signature=signature, input_example=input_example)
+        mlflow.keras.save_model(model, "my_model_mlflow", signature=signature, input_example=input_example)
\ No newline at end of file