diff --git a/predict/predict.py b/predict/predict.py
index fdeb1f7..02dc1f0 100644
--- a/predict/predict.py
+++ b/predict/predict.py
@@ -3,7 +3,7 @@ import mlflow
 import pandas as pd
 
 
-model = mlflow.sklearn.load_model("./prediction_data/movies_imdb")
+model = mlflow.keras.load_model("./prediction_data/movies_imdb")
 
 
 with open('./prediction_data/movies_imdb/input_example.json', 'r') as datafile: