Merge pull request #417 from pykale/modify_ckpt_in_cifa_cnntransformer

Add resume and test checkpoint loading to main.py in example/cifas_cn…

examples/cifar_cnntransformer/main.py

@@ -29,6 +29,8 @@ def arg_parse():
         help="gpu id(s) to use. int(0) for cpu. list[x,y] for xth, yth GPU."
         "str(x) for the first x GPUs. str(-1)/int(-1) for all available GPUs",
     )
+    parser.add_argument("--ckpt_resume", default="", help="path to train checkpoint file", type=str)
+    parser.add_argument("--ckpt_test", default="best", help="path to test checkpoint file", type=str)
     args = parser.parse_args()
     return args
 
@@ -82,10 +84,10 @@ def main():
     )
 
     # ---- start training ----
-    trainer.fit(model, train_loader, valid_loader)
+    trainer.fit(model, train_loader, valid_loader, ckpt_path=args.ckpt_resume)
 
     # ---- start testing ----
-    trainer.test(model, valid_loader)
+    trainer.test(model, valid_loader, ckpt_path=args.ckpt_test)
 
 
 if __name__ == "__main__":