From 8831f8e67bbae8deadbc845051d81f631581bcc2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20Wa=C5=82=C4=99sa?= Date: Sun, 15 May 2022 14:22:58 +0200 Subject: [PATCH] Zaktualizuj 'ml_pytorch_mlflow.py' --- ml_pytorch_mlflow.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ml_pytorch_mlflow.py b/ml_pytorch_mlflow.py index bf06f9e..29eef9b 100644 --- a/ml_pytorch_mlflow.py +++ b/ml_pytorch_mlflow.py @@ -148,17 +148,17 @@ def my_main(epochs): file.write(str(predict_single(input_, target, model))) input_example = val_ds[0] - signature = mlflow.models.signature.infer_signature(input_, prediction(input_, model)) + input_example = input_example.unsqueeze(0) + signature = mlflow.models.signature.infer_signature(input_, prediction(input_example, model)) tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme if tracking_url_type_store != "file": - mlflow.pytorch.log_model(model, "model", registered_model_name="s444356", signature=siganture, + mlflow.pytorch.log_model(model, "model", registered_model_name="s478839", signature=siganture, input_example=input_example) else: mlflow.pytorch.log_model(model, "model", signature=siganture, input_example=input_example) mlflow.pytorch.save_model(model, "my_model", signature=siganture, input_example=input_example) - torch.save(model, "Model_xPosition.pkl") with mlflow.start_run() as run: my_main(epochs)