1. 设置设备

  1. # 定义训练的设备
  2. device = torch.device("cuda" if torch.cuda.is_available else "cpu")

2. 将模型和优化器放到设备上

  1. model = Classify()
  2. # 将模型放到设备上
  3. model = model.to(device)
  4. loss_fn = nn.CrossEntropyLoss()
  5. # 将损失函数放到设备上
  6. loss_fn = loss_fn.to(device)

3. 将数据（imgs，targets）放到设备上

  1. # 将imgs和targets放到设备上去
  2. imgs = imgs.to(device)
  3. targets = targets.to(device)