pytorch保存模型遇到点问题

今天用pytorch保存模型时遇到bug

Can't pickle <class 'torch._C._VariableFunctions'>

在google上查找原因，发现是保存时保存了整个模型的原因，而模型中有一些自定义的参数

将torch.save(model,save_path)改为torch.save(model.state_dict(),save_path)

然后载入模型也做相应的更改就好了

参考链接https://github.com/pytorch/pytorch/issues/7545