From 55f61b9a708cf33e10071de353dac2638312d624 Mon Sep 17 00:00:00 2001 From: Yin Li Date: Sat, 12 Sep 2020 19:26:19 -0400 Subject: [PATCH] Add optional optimizer/scheduler loading --- map2map/train.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/map2map/train.py b/map2map/train.py index 7bf9c73..0898563 100644 --- a/map2map/train.py +++ b/map2map/train.py @@ -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