This commit is contained in:
Glenn Jocher 2019-09-19 18:08:21 +02:00
parent 5bacf9e0b8
commit de0612ca09
1 changed files with 5 additions and 5 deletions

View File

@ -108,11 +108,11 @@ def train():
chkpt = torch.load(weights, map_location=device)
# 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()}
model.load_state_dict(chkpt['model'], strict=False)
else:
model.load_state_dict(chkpt['model'])
# else:
# model.load_state_dict(chkpt['model'])
# load optimizer
if chkpt['optimizer'] is not None: