将参数加载到CPU
在使用torch.load()
方法加载函数的时候,会遇到CUDA OUT OF MEMORY的问题,这是由于训练的时候是在GPU上进行训练,因此在加载的时候默认也是加载到GPU上。
根据torch.load的官方文档torch.load('my_file.pt', map_location=lambda storage, loc: storage)
语句可以将参数保存到CPU上。
更多操作
1 | 'tensors.pt') torch.load( |
在使用torch.load()
方法加载函数的时候,会遇到CUDA OUT OF MEMORY的问题,这是由于训练的时候是在GPU上进行训练,因此在加载的时候默认也是加载到GPU上。
根据torch.load的官方文档torch.load('my_file.pt', map_location=lambda storage, loc: storage)
语句可以将参数保存到CPU上。
1 | >>> torch.load('tensors.pt') |