pytorch 中判断和指定模型和数据在GPU或CPU上

有时候，需要查看模型和数据是在GPU上，还是在CPU上；或者需要将模型和数据放在指定的GPU或者CPU上，那该怎么做呢？

• 1. 判断模型是在GPU上还是CPU上

model = nn.LSTM(input_size=10, hidden_size=4, num_layers=1, batch_first=True)
print(next(model.parameters()).device)  

• ２.判断数据是在GPU上还是CPU上

data = torch.ones([2, 3])
print(data.device) 

• ３.指定模型是在GPU上还是CPU上

model = nn.LSTM(input_size=10, hidden_size=4, num_layers=1, batch_first=True)
# 指定模型在GPU上
model = model.cuda()
# 指定模型在CPU上
model = model.cpu()

• 4.指定数据是在GPU上还是在CPU上

data = torch.ones([2, 3])
# 指定数据在GPU上 
data = data.cuda()
# 指定数据在CPU上 
data = data.cpu()

参考：https://www.cnblogs.com/picassooo/p/13736843.html