From 05066978fc4d18b628f36856c7b2d18249e2ff3d Mon Sep 17 00:00:00 2001
From: chenguowei01 <chenguowei01@baidu.com>
Date: Fri, 20 Dec 2019 17:21:04 +0800
Subject: [PATCH] eval for last epoch model

---
 pdseg/train.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pdseg/train.py b/pdseg/train.py
index 4ae5ad12..4f6a90e0 100644
--- a/pdseg/train.py
+++ b/pdseg/train.py
@@ -438,7 +438,8 @@ def train(cfg):
             except Exception as e:
                 print(e)
 
-        if epoch % cfg.TRAIN.SNAPSHOT_EPOCH == 0 and cfg.TRAINER_ID == 0:
+        if (epoch % cfg.TRAIN.SNAPSHOT_EPOCH == 0
+                or epoch == cfg.SOLVER.NUM_EPOCHS) and cfg.TRAINER_ID == 0:
             ckpt_dir = save_checkpoint(exe, train_prog, epoch)
 
             if args.do_eval:
-- 
GitLab