torch用GPU训练

需要手动将数据和模型都放到GPU中

法一

	# 是否用GPU训练
	train_on_gpu = torch.cuda.is_available()
	device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
	model = model.to(device)
	data = data.to(device)

法二

model = model.cuda()
data = data.cuda()
#还是法二好用

判断自己的模型或者数据在哪儿

# 1. 通过判断模型model的参数是否在cuda上来判定模型是否在gpu上。

print('Is model on gpu: ', next(model.parameters()).is_cuda)
输出若是True，则model在gpu上

2. 输出数据data的device字段。

print('data device: ', data.device)