From 31201b169b76b61c5f2dcdb02c0703cee452ab60 Mon Sep 17 00:00:00 2001
From: Wojciech Jarmosz <wojciech@jarmosz.com>
Date: Sun, 23 May 2021 16:57:24 +0200
Subject: [PATCH] Change flavor of model prediction

---
 mlflow_prediction.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlflow_prediction.py b/mlflow_prediction.py
index 8ddb490..958358b 100644
--- a/mlflow_prediction.py
+++ b/mlflow_prediction.py
@@ -9,7 +9,7 @@ arguments = sys.argv[1:]
 
 input = str(arguments[0])
 
-model = mlflow.keras.load_model("my_model")
+model = mlflow.pytorchg.load_model("my_model")
 
 with open('my_model/{input}', 'r') as file:
     json_data = json.load(file)['inputs']