diff --git a/avocado-mlflow.py b/avocado-mlflow.py
index cdbd29a..ec2649b 100644
--- a/avocado-mlflow.py
+++ b/avocado-mlflow.py
@@ -77,6 +77,6 @@ with mlflow.start_run():
     mlflow.log_param("batch_size", batch_size)
     mlflow.log_metric("rmse", rmse) 
     #mlflow.keras.log_model(model, 'avocado_model.h5')
-    mlflow.keras.log_model(model, 'avocado_model', registered_model_name="s434742", signature=infer_signature(X_train, y_train), input_example=X_train.iloc[0])
-    mlflow.keras.save_model(model, 'avocado_model', signature=infer_signature(X_train, y_train), input_example=X_train.iloc[0])
+    mlflow.keras.log_model(model, 'avocado_model', registered_model_name="s434742", signature=infer_signature(X_train, y_train), input_example=np.array(X_train.iloc[0]))
+    mlflow.keras.save_model(model, 'avocado_model', signature=infer_signature(X_train, y_train), input_example=np.array(X_train.iloc[0]))