Update app.py

app.py

@@ -10,13 +10,10 @@ class Prediction:
     prediction: TorchTensor
 
 app = fastapi.FastAPI(docs_url="/")
-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']
+# Load the pre-trained model
+pre_trained_model = torch.load('best_model.pth', map_location=torch.device('cpu'))
 
-# Load the state dictionary into the current model
-model = pre_trained_model
 # Define a function to preprocess the input image
 def preprocess_input(input: fastapi.UploadFile):
     image = Image.open(input.file)