Fix: Use GPU

app.py

@@ -21,7 +21,7 @@ def load_model(model_path: str):
     model = models.resnet50(weights=None)
 
     # Load custom weights from a .pth file with CPU mapping
-    state_dict = torch.load(model_path, map_location=torch.device('cpu'))
+    state_dict = torch.load(model_path)
 
     # 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()}