Skip to content
Snippets Groups Projects
Commit 09d2eea3 authored by Neta Zmora's avatar Neta Zmora
Browse files

Allow loading checkpoints not containing an optimizer

Support loading a model from a checkpoint file that does not
have an Optimizer instance.
Before the change, loading such a model required using
```load_lean_checkpoint```  (or --exp-load-weights-from
From the compress_classifier.py command-line), therefore
this change is for convenience only.
parent a69dd5d6
No related branches found
No related tags found
No related merge requests found
...@@ -185,10 +185,9 @@ def load_checkpoint(model, chkpt_file, optimizer=None, model_device=None, *, lea ...@@ -185,10 +185,9 @@ def load_checkpoint(model, chkpt_file, optimizer=None, model_device=None, *, lea
optimizer = _load_optimizer(checkpoint['optimizer_type'], optimizer = _load_optimizer(checkpoint['optimizer_type'],
checkpoint['optimizer_state_dict'], model) checkpoint['optimizer_state_dict'], model)
except KeyError: except KeyError:
if 'optimizer' not in checkpoint: # Older checkpoints do support optimizer loading: They either had an 'optimizer' field
raise # (different name) which was not used during the load, or they didn't even checkpoint
# older checkpoints didn't support this feature # the optimizer.
# they had the 'optimizer' field instead
optimizer = None optimizer = None
if optimizer is not None: if optimizer is not None:
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment