多卡训练转换到单卡训练,在保存模型时发生了错误:
torch.nn.modules.module.ModuleAttributeError: 'VGG' object has no attribute 'module'
后来发现是保存模型的时候用了torch.module.statedict()
stateG = {'model': netG.module.cpu().state_dict(), 'optimizer': optimizerG.state_dict(), 'epoch': epoch}
torch.save(stateG, net_g_model_out_path)
上网查找资料后将代码改为:
stateG = {'model': netG.cpu().state_dict(), 'optimizer': optimizerG.state_dict(), 'epoch': epoch}
torch.save(stateG, net_g_model_out_path)