Update prediction.py

prediction.py

@@ -15,7 +15,7 @@ class Prediction:
         Initialize the Prediction class with a pre-trained model and necessary parameters.
         """
         self.device = torch.device("cpu")
-        self.model = torch.jit.load("model.pt")
+        self.model = torch.jit.load("model.pt", map_location=self.device)
 
         self.model.eval()
         params = read_yaml(PARAMS_FILE_PATH)