bug-pytorch

torch 指定显卡的常见方式有一下两种

 

第一：
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '1'

第二：
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')

　　

区别在与第一种一开始会将初始化的模型，以及模型和数据放入显卡1当中，而第二种的话，初始化模型仍然在显卡0当中，当显卡0被完全占满时，则会爆内存错误

 

将模型完全放在cpu上

torch.load(model_file_path,map_location='cpu')