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

import torch
import torch.nn as nn

# ----------- 判断模型是否在GPU ----------------------

model = nn.LSTM(input_size=10, hidden_size=4, num_layers=1, batch_first=True)
print(next(model.parameters()).is_cuda)  # 输出：False

model = model.cuda()
print(next(model.parameters()).is_cuda)  # 输出：True

model = model.cpu()
print(next(model.parameters()).is_cuda)  # 输出：False

# ----------- 判断数据是在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