diff --git a/.gitignore b/.gitignore
index 1d21078..0208904 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,4 +8,6 @@ my_runs
 mlruns
 my_model
 1/
-mydb.sqlite
\ No newline at end of file
+mydb.sqlite
+movies_on_streaming_platforms_model.zip
+movies_on_streaming_platforms_model
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
index 0433de2..27031a0 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -15,6 +15,7 @@ RUN pip3 install torch==1.8.1+cpu torchvision==0.9.1+cpu torchaudio==0.8.1 -f ht
 RUN pip3 install sacred
 RUN pip3 install pymongo
 RUN pip3 install mlflow
+RUN pip3 install tensorflow==2.5.0rc1
 # Stwórzmy w kontenerze (jeśli nie istnieje) katalog /app i przejdźmy do niego (wszystkie kolejne polecenia RUN, CMD, ENTRYPOINT, COPY i ADD będą w nim wykonywane)
 WORKDIR /app
 
diff --git a/mlflow_predict.py b/mlflow_predict.py
index 3f290b4..6a8a4f9 100644
--- a/mlflow_predict.py
+++ b/mlflow_predict.py
@@ -1,17 +1,14 @@
 import mlflow
-import mlflow.pytorch
+import mlflow.keras
 import sys
 import json
-import numpy as np
-import torch
 
 input = sys.argv[1]
 
-model = mlflow.pytorch.load_model("my_model")
+model = mlflow.keras.load_model("movies_on_streaming_platforms_model")
 
-with open('my_model/'+input) as json_file:
+with open('movies_on_streaming_platforms_model/'+input) as json_file:
     data = json.load(json_file)
-#print(np.array(data['inputs']))
-print(model(torch.tensor(np.array(data['inputs'])).float()))
-
+#print(data)
+print(model.predict(data['inputs']))
             
\ No newline at end of file
diff --git a/mlflow_predict_registry.py b/mlflow_predict_registry.py
index 47339a5..2f84527 100644
--- a/mlflow_predict_registry.py
+++ b/mlflow_predict_registry.py
@@ -1,16 +1,14 @@
-import mlflow
-import mlflow.pytorch
 from mlflow.tracking import MlflowClient
-import numpy as np
-import torch
+import mlflow
+import mlflow.keras
 import json
 
 #mlflow.set_tracking_uri("http://127.0.0.1:5000")
 mlflow.set_tracking_uri("http://172.17.0.1:5000")
 client = MlflowClient()
 version = 0
-model_name = "s426206"
-for mv in client.search_model_versions("name='s426206'"):
+model_name = "s434704"
+for mv in client.search_model_versions(f"name='{model_name}'"):
     if int(mv.version) > version:
         version = int(mv.version)
 
@@ -18,7 +16,7 @@ model = mlflow.pytorch.load_model(
     model_uri=f"models:/{model_name}/{version}"
 )
 
-with open('my_model/input_example.json') as json_file:
+with open('movies_on_streaming_platforms_model/input_example.json') as json_file:
     data = json.load(json_file)
 #print(np.array(data['inputs']))
-print(model(torch.tensor(np.array(data['inputs'])).float()))
\ No newline at end of file
+print(model.predict(data['inputs']))
\ No newline at end of file