From 65857adf52f8936dfd36c1ad07ff75f3be14d813 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 18 Jul 2020 14:34:12 -0700 Subject: [PATCH] update train.py ckpt loading --- train.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/train.py b/train.py index 879bb2e..c88e3b6 100644 --- a/train.py +++ b/train.py @@ -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. " \