Update app.py

app.py

@@ -36,8 +36,9 @@ class CustomEfficientNet(nn.Module):
 
 # Définir le modèle final
 num_classes = 2
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = CustomEfficientNet(base_model, num_classes).to("cuda" if torch.cuda.is_available() else "cpu")
-model.load_state_dict(torch.load("efficientnet_b7_best.pth",weights_only=False))
+model.load_state_dict(torch.load("efficientnet_b7_best.pth",weights_only=False, map_location=device))
 model.eval()  # Passer le modèle en mode évaluation