[fix] replace ema_model with model in demo/test_ap_on_coco (#242)

2023-11-13 13:52:37 +08:00 · 2023-11-13 13:52:37 +08:00 · eeba084341
parent 60d796825e
commit eeba084341
1 changed files with 1 additions and 1 deletions
--- a/demo/test_ap_on_coco.py
+++ b/demo/test_ap_on_coco.py
@ -26,7 +26,7 @@ def load_model(model_config_path: str, model_checkpoint_path: str, device: str =
    args.device = device
    model = build_model(args)
    checkpoint = torch.load(model_checkpoint_path, map_location="cpu")
-    model.load_state_dict(clean_state_dict(checkpoint["ema_model"]), strict=False)
+    model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False)
    model.eval()
    return model