From af17a9a8ca44f19aec3d1d566c839091b07058cf Mon Sep 17 00:00:00 2001 From: Abhay Deshpande Date: Fri, 9 Dec 2022 02:06:49 +0000 Subject: [PATCH] Move key declaratin out of if --- src_files/models/utils/factory.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src_files/models/utils/factory.py b/src_files/models/utils/factory.py index 12a2f34..8c2aea5 100644 --- a/src_files/models/utils/factory.py +++ b/src_files/models/utils/factory.py @@ -44,11 +44,11 @@ def create_model(args,load_head=False): model_path = "./tresnet_l.pth" print('done') state = torch.load(model_path, map_location='cpu') + if 'model' in state: + key = 'model' + else: + key = 'state_dict' if not load_head: - if 'model' in state: - key = 'model' - else: - key = 'state_dict' filtered_dict = {k: v for k, v in state[key].items() if (k in model.state_dict() and 'head.fc' not in k)} model.load_state_dict(filtered_dict, strict=False)