Update app.py

app.py

@@ -10,7 +10,7 @@ class Prediction:
     prediction: TorchTensor
 
 app = fastapi.FastAPI(docs_url="/")
-pre_trained_model = torch.load('best_model.pth')
+pre_trained_model = torch.load('best_model.pth',map_location=torch.device('cpu'))
 
 # Get the state dictionary of the pre-trained model
 state_dict = pre_trained_model['model_state_dict']