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

@ -105,6 +105,8 @@ def train():
if opt.transfer:
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)
else:
model.load_state_dict(chkpt['model'])
# load optimizer
if chkpt['optimizer'] is not None: