torch.nn.modules.module.ModuleAttributeError: 'MobileNetV2' object has no attribute 'module'

[epoc88]

This happens when try to use test.py to load the model:

pfld_backbone.load_state_dict(torch.load(args.model_path, map_location= device))

To fix this issue,
we have to change save model in the "train_model.py"

---

save model

checkpoint_path = os.path.join(model_dir, 'model_'+str(epoch)+'.pth')
if args.all_model:
torch.save(model, checkpoint_path)

as
torch.save(model.state_dict(), checkpoint_path)

after that all works.

I am also wondering about an alternative without changing the train_model.py. So instead of using the train_model.py, we could do the code like below:

model = torch.load(pretrained_model)
test(test_loader, model, args, device)
...