set learning rate for restored model

This commit is contained in:
Eren Golge 2019-02-12 10:04:39 +01:00
parent d03865ed05
commit 4fad1971cd
1 changed files with 2 additions and 0 deletions

View File

@ -389,6 +389,8 @@ def main(args):
criterion.cuda() criterion.cuda()
criterion_st.cuda() criterion_st.cuda()
optimizer.load_state_dict(checkpoint['optimizer']) optimizer.load_state_dict(checkpoint['optimizer'])
for group in optimizer.param_groups:
group['lr'] = c.lr
print( print(
" > Model restored from step %d" % checkpoint['step'], flush=True) " > Model restored from step %d" % checkpoint['step'], flush=True)
start_epoch = checkpoint['epoch'] start_epoch = checkpoint['epoch']