In the demo code, update the lines to load the model as below:

[ihongxx]

In the demo code, update the lines to load the model as below:

model = get_model('vgg19')
wts_dict = torch.load(args.weight)
wts_dict_corrected = {k.replace('module.',''): v for k, v in wts_dict.items()}
model.load_state_dict(wts_dict_corrected)

Originally posted by @SaraSherinThomas in #145 (comment)