pytorch 如何在GPU上训练

net = AlexNet()
net.cuda()#转移到CUDA上

criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()

loss = criterion(out, label)

#download the dataset
train_set = CIFAR10("./data_cifar10", train=True, transfORM=data_tf, download=True)
train_data = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)

    for im, label in train_data:
        i = i + 1
        im = im.cuda()#把数据迁移到CUDA上
        im = Variable(im)#把数据放到Variable里
        label = label.cuda()
        label =Variable(label)
        out = net(im)#the output should have the size of (N,10)

nvidia -smi

watch -n -1 nvidia -smi

import os
# 使用第一张与第三张GPU卡
os.environ["CUDA_VISIBLE_DEVICES"] = "0,3"