Update app.py

app.py

@@ -10,7 +10,7 @@ model_file = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
 
 model = UNet()
 state = torch.load(model_file, map_location="cpu")
-model.load_state_dict(state)
+model.load_state_dict(state, strict=False)  # use strict=False to ignore mismatched keys
 model.eval()
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)