Update app.py

app.py

@@ -14,7 +14,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load model
 model = UNet().to(device)  # Move the model to the selected device
-model_state_dict = torch.load(r"model.pth", map_location=device)
+model_state_dict = torch.load(r"model.pth", map_location=device, weights_only=True)
 model.load_state_dict(model_state_dict["model_state_dict"])
 
 scaler = min_max_scaler()