Add optional optimizer/scheduler loading

2020-09-12 19:26:19 -04:00 · 2020-09-12 19:26:19 -04:00 · 55f61b9a70
commit 55f61b9a70
parent b215c0fa7b
1 changed files with 4 additions and 2 deletions
--- a/map2map/train.py
+++ b/map2map/train.py
@ -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