diff --git a/avocado-mlflow.py b/avocado-mlflow.py
index c3d7763..71197f2 100644
--- a/avocado-mlflow.py
+++ b/avocado-mlflow.py
@@ -1,5 +1,6 @@
 import sys
 from keras.backend import batch_dot, mean
+from mlflow.models.signature import infer_signature
 import pandas as pd
 import numpy as np
 from six import int2byte
@@ -53,7 +54,8 @@ def my_main(epochs, batch_size):
     # zapisanie modelu
     model.save('avocado_model.h5')
 
-    return rmse, model
+
+    return rmse, model, X_train, y_train
 
 
 
@@ -63,9 +65,10 @@ batch_size = int(sys.argv[2]) if len(sys.argv) > 2 else 16
 
 with mlflow.start_run():
         
-    rmse, model = my_main(epochs, batch_size)
+    rmse, model, X_train, y_train = my_main(epochs, batch_size)
 
     mlflow.log_param("epochs", epochs)
     mlflow.log_param("batch_size", batch_size)
     mlflow.log_metric("rmse", rmse) 
-    mlflow.keras.log_model(model, 'avocado_model.h5')
\ No newline at end of file
+    #mlflow.keras.log_model(model, 'avocado_model.h5')
+    mlflow.keras.log_model(keras_model=model, path='avocado_model', signature=infer_signature(X_train, y_train), input_example=X_train.iloc[0])
\ No newline at end of file