Aquila-Consortium/map2map
32
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add optional optimizer/scheduler loading

Browse Source
This commit is contained in:
Yin Li 2020-09-12 19:26:19 -04:00
parent b215c0fa7b
commit 55f61b9a70
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
map2map/train.py
View File

@ -159,8 +159,10 @@ def gpu_worker(local_rank, node, args):
load_model_state_dict(model.module, state['model'],
strict=args.load_state_strict)
optimizer.load_state_dict(state['optimizer'])
scheduler.load_state_dict(state['scheduler'])
if 'optimizer' in state:
optimizer.load_state_dict(state['optimizer'])
if 'scheduler' in state:
scheduler.load_state_dict(state['scheduler'])
torch.set_rng_state(state['rng'].cpu()) # move rng state back