PyTorch查看模型和数据是否在GPU上

 转载自：https://www.cnblogs.com/picassooo/p/13736843.html，本文只做个人记录学习使用，版权归原作者所有。

import torch
import torch.nn as nn
 
# ----------- 判断模型是在CPU还是GPU上 ----------------------
 
model = nn.LSTM(input_size=10, hidden_size=4, num_layers=1, batch_first=True)
print(next(model.parameters()).device)  # 输出：cpu
 
model = model.cuda()
print(next(model.parameters()).device)  # 输出：cuda:0
 
model = model.cpu()
print(next(model.parameters()).device)  # 输出：cpu
 
# ----------- 判断数据是在CPU还是GPU上 ----------------------
 
data = torch.ones([2, 3])
print(data.device)  # 输出：cpu
 
data = data.cuda()
print(data.device)  # 输出：cuda:0
 
data = data.cpu()
print(data.device)  # 输出：cpu

此外，用.is_cuda也可以判断模型和数据是否在GPU上，例如： data.is_cuda