GPU训练

单个GPU进行训练

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
inputs, labels = inputs.to(device), labels.to(device)

多个GPU进行训练

MULTI-GPU EXAMPLES
DATA PARALLELISM

报错

Expected object of device type cuda but got device type cpu for argument #1

1 模型是否放到了CUDA上 model = model.to(device) 或 model = model.cuda(device)
2 输入数据是否放到了CUDA上 data = data.to(device) 或 data = data .cuda(device)
3 模型内部新建的张量是否放到了CUDA上 p = torch.tensor([1]).cuda(device) 或 p = torch.tensor([1]).cuda(device)