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)