PyTorch查看模型和數據是否在GPU上

模型和數據可以在CPU和GPU上來回遷移，怎么判斷模型和數據在哪里呢？

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