From 45ee10ccc0fd3377aede1fb940517b417c94f1b2 Mon Sep 17 00:00:00 2001 From: Filip Izydorczyk Date: Fri, 4 Jun 2021 16:10:51 +0200 Subject: [PATCH] mlflow pytorch lib update --- learning/ml-mlflow.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/learning/ml-mlflow.py b/learning/ml-mlflow.py index 70e8298..291dcbf 100644 --- a/learning/ml-mlflow.py +++ b/learning/ml-mlflow.py @@ -80,10 +80,10 @@ with mlflow.start_run(): mlflow.set_experiment("s434700") tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme if tracking_url_type_store != "file": - mlflow.keras.log_model(model, "model.pt", registered_model_name="s434700", signature=signature, - input_example=test_input) + mlflow.pytorch.log_model(model, "model.pt", registered_model_name="s434700", signature=signature, + input_example=test_input) else: - mlflow.keras.log_model(model, "model.pt", - signature=signature, input_example=test_input) - mlflow.keras.save_model( + mlflow.pytorch.log_model(model, "model.pt", + signature=signature, input_example=test_input) + mlflow.pytorch.save_model( model, "model.pt", signature=signature, input_example=test_input)