From 28dbfa5d6f59b5c06e556545f6d6e141eedf30f7 Mon Sep 17 00:00:00 2001 From: Adrian Charkiewicz Date: Sun, 15 May 2022 20:13:30 +0200 Subject: [PATCH] signature should by numpy array --- pytorch/pytorch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch/pytorch.py b/pytorch/pytorch.py index 6c537d2..5dbd84b 100644 --- a/pytorch/pytorch.py +++ b/pytorch/pytorch.py @@ -186,7 +186,7 @@ def fit(epochs, lr, model, train_loader, val_loader, _log, _run, opt_func=torch. _run.info["epochs"] = epochs - signature = mlflow.models.signature.infer_signature(train_ds) + signature = mlflow.models.signature.infer_signature(inputs_array) tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme