From 4427067c0b6e6d70b34abfb3fddbc5b2808bb9f7 Mon Sep 17 00:00:00 2001 From: Eren Date: Sat, 29 Sep 2018 01:45:33 +0200 Subject: [PATCH] restore lr scheduler --- train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train.py b/train.py index ffe2557e..a43754d3 100644 --- a/train.py +++ b/train.py @@ -429,7 +429,7 @@ def main(args): criterion.cuda() criterion_st.cuda() - scheduler = AnnealLR(optimizer, warmup_steps=c.warmup_steps) + scheduler = AnnealLR(optimizer, warmup_steps=c.warmup_steps, last_epoch=args.restore_step) num_params = count_parameters(model) print(" | > Model has {} parameters".format(num_params), flush=True)