inconsistent implementation for saved .pth model about hrnet-facial-landmark-detection HOT 2 OPEN

@JMingKuo, I modify the keys in the test.py as your suggestion. But there is still an error. Is there something wrong when I modify the code? I have used your code instead of the code in test.py as follow.
Thanks!

state_dict = torch.load(args.model_file)
if 'state_dict' in state_dict.keys():
state_dict = state_dict['state_dict']
model.load_state_dict(state_dict)
else:
model.module.load_state_dict(state_dict)

from hrnet-facial-landmark-detection.

@briantse100, you only have to modify the keys while using 300W model and always use model.load_state_dict for loading weights.

A more complete modification would like:

state_dict = torch.load(args.model_file)
if USING_300W:
    new_state_dict = OrderedDict()
    for key in state_dict:
        new_state_dict[key.replace('module.', '')] = state_dict[key]
    model.load_state_dict(new_state_dict)
else:
    model.load_state_dict(state_dict)

from hrnet-facial-landmark-detection.