diff --git a/mlflow_prediction.py b/mlflow_prediction.py
index 75791a2..22abfc5 100644
--- a/mlflow_prediction.py
+++ b/mlflow_prediction.py
@@ -14,6 +14,9 @@ input = str(arguments[0])
 model = mlflow.pytorch.load_model("my_model")
 
 with open(f'my_model/{input}', 'r') as file:
-    json_data = json.load(file)['inputs']
+    json_data = json.load(file)
+
+print(json_data)
+print(json_data['inputs'])
 
 print(model(torch.tensor(np.array(json_data['inputs'])).float()))