set learning rate for restored model

2019-02-12 10:04:39 +01:00 · 2019-02-12 10:04:39 +01:00 · ca5b7d6c69
parent d28bbe09fb
commit ca5b7d6c69
1 changed files with 2 additions and 0 deletions
--- a/train.py
+++ b/train.py
@ -389,6 +389,8 @@ def main(args):
            criterion.cuda()
            criterion_st.cuda()
        optimizer.load_state_dict(checkpoint['optimizer'])
+        for group in optimizer.param_groups:
+                group['lr'] = c.lr
        print(
            " > Model restored from step %d" % checkpoint['step'], flush=True)
        start_epoch = checkpoint['epoch']