You are saving the nn.DataParallel state_dict, which will include the .module keys as described here 4. You could either wrap the model first into nn.DataParallel and load the state_dict afterwards or store the net.module.state_dict() instead.