Add optional optimizer/scheduler loading

This commit is contained in:
Yin Li 2020-09-12 19:26:19 -04:00
parent b215c0fa7b
commit 55f61b9a70

View File

@ -159,7 +159,9 @@ def gpu_worker(local_rank, node, args):
load_model_state_dict(model.module, state['model'],
strict=args.load_state_strict)
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