[Pytorch学习]在GPU上训练

Pytorch学习之在GPU上训练

主要分为以下:

  1. 定义GPU
device=torch.device("cuda:0"if torch.cuda.is_available() else "cpu")
  1. 将模型放到GPU中
def Model(...):
	....
	....

model=Model().to(device)
  1. 把输入和标签放到GPU中
inputs, labels= inputs.to(device), labels.to(device)
  1. 接下来就可以正常训练拉

猜你喜欢

转载自blog.csdn.net/crabstew/article/details/88996645