Fix: Runtime error for CPU

app.py

@@ -20,14 +20,15 @@ def load_model(model_path: str):
     # Load the pre-trained ResNet50 model from ImageNet
     model = models.resnet50(pretrained=False)
 
-    # Load custom weights from a .pth file
-    state_dict = torch.load(model_path)
+    # Load custom weights from a .pth file with CPU mapping
+    state_dict = torch.load(model_path, map_location=torch.device('cpu'))
 
     # Filter out unexpected keys
     filtered_state_dict = {k: v for k, v in state_dict['model_state_dict'].items() if k in model.state_dict()}
 
     # Load the filtered state dictionary into the model
-    model.load_state_dict(filtered_state_dict, strict=False) 
+    model.load_state_dict(filtered_state_dict, strict=False)
+    model.eval()
     return model