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)
- 接下来就可以正常训练拉
版权声明:本文为crabstew原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接和本声明。