提交 65857adf authored 作者: Glenn Jocher's avatar Glenn Jocher

update train.py ckpt loading

上级 7b0c9be8
......@@ -119,8 +119,7 @@ def train(hyp):
# load model
try:
ckpt['model'] = {k: v for k, v in ckpt['model'].float().state_dict().items()
if model.state_dict()[k].shape == v.shape} # to FP32, filter
ckpt['model'] = {k: v for k, v in ckpt['model'].float().state_dict().items() if k in model.state_dict()}
model.load_state_dict(ckpt['model'], strict=False)
except KeyError as e:
s = "%s is not compatible with %s. This may be due to model differences or %s may be out of date. " \
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论