diff --git a/train-mlflow.py b/train-mlflow.py
index 5333775..8a59205 100644
--- a/train-mlflow.py
+++ b/train-mlflow.py
@@ -107,7 +107,7 @@ def my_main(epochs):
         
 
             input_example = X_testing[0]
-            siganture = mlflow.models.signature.infer_signature(X_training.eval(), y_predicted_train.eval())
+            siganture = mlflow.models.signature.infer_signature(np.array(X_training), np.array(y_predicted_train))
             tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
 
             if tracking_url_type_store != "file":