Pytorch学习之在GPU上训练
主要分为以下:
- 定义GPU
device=torch.device("cuda:0"if torch.cuda.is_available() else "cpu")
- 将模型放到GPU中
def Model(...):
....
....
model=Model().to(device)
- 把输入和标签放到GPU中
inputs, labels= inputs.to(device), labels.to(device)
- 接下来就可以正常训练拉