allow test for checkpoint from cfg.MODEL.WEIGHT (facebookresearch#894)

Lyears · Jun 14, 2019 · 6bbaf85 · 6bbaf85
1 parent 01c3d0b
commit 6bbaf85
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tools/test_net.py b/tools/test_net.py
@@ -78,7 +78,8 @@ def main():
 
     output_dir = cfg.OUTPUT_DIR
     checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir)
-    _ = checkpointer.load(args.ckpt, use_latest=args.ckpt is None)
+    ckpt = cfg.MODEL.WEIGHT if args.ckpt is None else args.ckpt
+    _ = checkpointer.load(ckpt, use_latest=args.ckpt is None)
 
     iou_types = ("bbox",)
     if cfg.MODEL.MASK_ON: