Update app.py

app.py

@@ -34,13 +34,16 @@ with open(config, 'r') as f:
     config = yaml.load(f, Loader=yaml.FullLoader)
 hf_weight = hf_hub_download(repo_id=f"ZidongC/PanDA", filename=f"panda_large.pth", repo_type="model")
 state_dict = torch.load(hf_weight, map_location="cpu")
+new_state_dict = {}
+for key, value in state_dict.items():
+    new_key = key[7:] if key.startswith('module.') else key
+    new_state_dict[new_key] = value
 model = make(config['model'])
-if any(key.startswith('module') for key in state_dict.keys()):
-    model = nn.DataParallel(model)
+# if any(key.startswith('module') for key in state_dict.keys()):
+#     model = nn.DataParallel(model)
 model_state_dict = model.state_dict()
-model.load_state_dict({k: v for k, v in state_dict.items() if k in model_state_dict})
+model.load_state_dict({k: v for k, v in new_state_dict.items() if k in model_state_dict})
 model = model.to(DEVICE).eval()
-model = model.module
 
 title = "# PanDA"
 description = """Official demo for **PanDA**.