restore lr scheduler

2018-09-29 01:45:33 +02:00 · 2018-09-29 01:45:33 +02:00 · 4427067c0b
parent 3d3fb14a4b
commit 4427067c0b
1 changed files with 1 additions and 1 deletions
--- 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)