From 50599e264a91c3330b4349e2a99d6485360b012c Mon Sep 17 00:00:00 2001 From: wangguanzhong Date: Wed, 25 Nov 2020 13:16:29 +0800 Subject: [PATCH] fix model save for last epoch (#1761) --- tools/train.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tools/train.py b/tools/train.py index 318ae855f..fd4c3694d 100755 --- a/tools/train.py +++ b/tools/train.py @@ -196,7 +196,9 @@ def run(FLAGS, cfg, place): logger.info(strs) # Save Stage - if ParallelEnv().local_rank == 0 and cur_eid % cfg.snapshot_epoch == 0: + if ParallelEnv().local_rank == 0 and ( + cur_eid % cfg.snapshot_epoch == 0 or + (cur_eid + 1) == int(cfg.epoch)): cfg_name = os.path.basename(FLAGS.config).split('.')[0] save_name = str(cur_eid) if cur_eid + 1 != int( cfg.epoch) else "model_final" -- GitLab