diff --git a/predict_s444356.py b/predict_s444356.py
index 6060158..9879b5d 100644
--- a/predict_s444356.py
+++ b/predict_s444356.py
@@ -1,8 +1,6 @@
 import json
 import mlflow
 import numpy as np
-import torch
-from torch.autograd import Variable
 
 logged_model = 'mlruns/1/d5b6f9c1784a4d2dbb8592cd4ad364d7/artifacts/model'
 loaded_model = mlflow.pyfunc.load_model(logged_model)
@@ -10,7 +8,7 @@ loaded_model = mlflow.pyfunc.load_model(logged_model)
 
 with open(f'{logged_model}/input_example.json') as f:
     data = json.load(f)
-    input_example = np.array(data['inputs'][0])
-    input_example = Variable(torch.from_numpy(input_example)).float()
+    input_example = np.array([data['inputs'][0]], dtype=np.float64).reshape(-1, 2)
 
-loaded_model.predict(input_example)
+
+print(f'Prediction: {loaded_model.predict(input_example)}')