cuda runtime error (2) : out of memory at /pytorch/torch/csrc/generic/serialization.cpp:15
应该是一个bug, 参考链接,参考评论,将模型转换到CPU之后,保存,再将模型转换到GPU即可以解决问题。
代码如下:
atten_t_net = atten_t_net.to(torch.device('cpu'))
torch.save(atten_t_net.state_dict(), pkl_name)
atten_t_net = atten_t_net.to('cuda:0')