restore lr scheduler

2026-09-09 11:16:00 +08:00 · 2018-09-29 01:46:11 +02:00
parent 3d3fb14a4b
commit 4427067c0b
1 changed files with 1 additions and 1 deletions
@@ -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)