Solved by ptrblck in post #2 
              
              
                This error will be raised, if you’ve stored the state_dict from the nn.DataParallel model and are now trying to load it to a plain model. 
You could use this approach to remove the .module from the keys or alternatively store the state_dict via: 
sd = model.module.state_dict()
torch.save(sd, PATH)