diff --git a/learning/ml-mlflow.py b/learning/ml-mlflow.py
index 116a085..70e8298 100644
--- a/learning/ml-mlflow.py
+++ b/learning/ml-mlflow.py
@@ -75,7 +75,7 @@ with mlflow.start_run():
     predicted = model(Variable(torch.from_numpy(x_train))).data.numpy()
 
     signature = mlflow.models.signature.infer_signature(
-        x_train.values, predicted)
+        x_train, predicted)
 
     mlflow.set_experiment("s434700")
     tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme