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']