Merge pull request #5 from carstenli/patch-3
build_backbone.py update to make it work on Apple Silicon devices
This commit is contained in:
@@ -26,7 +26,7 @@ def build_backbone(bb_name, pretrained=True, params_settings=''):
|
||||
return bb
|
||||
|
||||
def load_weights(model, model_name):
|
||||
save_model = torch.load(config.weights[model_name])
|
||||
save_model = torch.load(config.weights[model_name], map_location=torch.device('cpu'))
|
||||
model_dict = model.state_dict()
|
||||
state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model.items() if k in model_dict.keys()}
|
||||
# to ignore the weights with mismatched size when I modify the backbone itself.
|
||||
@@ -35,7 +35,7 @@ def load_weights(model, model_name):
|
||||
sub_item = save_model_keys[0] if len(save_model_keys) == 1 else None
|
||||
state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model[sub_item].items() if k in model_dict.keys()}
|
||||
if not state_dict or not sub_item:
|
||||
print('Weights are not successully loaded. Check the state dict of weights file.')
|
||||
print('Weights are not successfully loaded. Check the state dict of weights file.')
|
||||
return None
|
||||
else:
|
||||
print('Found correct weights in the "{}" item of loaded state_dict.'.format(sub_item))
|
||||
|
||||
Reference in New Issue
Block a user