gpu训练的两种方式

1.训练模型，数据（imgs，targets），loss函数后面加.cuda来进行gpu训练（训练模型和loss可以不用另外赋值，直接.cuda()，但是数据必须要进行赋值）

tudui = Tudui()#模型定义
if torch.cuda.is_available():
    tudui = tudui.cuda()
#loss函数
loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
    loss_fn = loss_fn.cuda()
#imgs和targets
imgs,targets = data
if torch.cuda.is_available():
    imgs.cuda()
    targets.cuda()

2.在训练模型，数据，loss后.to(device)，提前定义torch.device()("cpu","cuda:0","cuda:1")，如下一个小栗子

#定义训练的设备
device = torch.device("cpu")
tudui = Tudui()
tudui = tudui.to(device)