This commit is contained in:
Glenn Jocher 2019-08-23 15:27:29 +02:00
parent 356c85bf0e
commit 8ef49f2560
1 changed files with 3 additions and 1 deletions

View File

@ -104,7 +104,9 @@ def train():
# load model # load model
if opt.transfer: if opt.transfer:
chkpt['model'] = {k: v for k, v in chkpt['model'].items() if model.state_dict()[k].numel() == v.numel()} chkpt['model'] = {k: v for k, v in chkpt['model'].items() if model.state_dict()[k].numel() == v.numel()}
model.load_state_dict(chkpt['model'], strict=False) model.load_state_dict(chkpt['model'], strict=False)
else:
model.load_state_dict(chkpt['model'])
# load optimizer # load optimizer
if chkpt['optimizer'] is not None: if chkpt['optimizer'] is not None: