解决Expected all tensors to be on the same device, but found at least two devices, cuda:0
(1)可能是模型没有移动到和数据相同的device上:`model.to(device)`(2)可能是input和参数没有在相同device上
·
一、问题描述
在跑pytorch代码时报错:
Expected all tensors to be on the same device, but found at least two devices, cuda:0
二、解决方案
判断参数数据是在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
查看模型的参数是在哪里(其实也可以直接data.is_cuda
):
model = nn.LSTM(input_size=10, hidden_size=4, num_layers=1, batch_first=True)
print(next(model.parameters()).device)
(1)可能是模型没有移动到和数据相同的device上:model.to(device)
(2)可能是input和参数没有在相同device上,如下类似的操作即可:
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
data.to(device)
欢迎来到由智源人工智能研究院发起的Triton中文社区,这里是一个汇聚了AI开发者、数据科学家、机器学习爱好者以及业界专家的活力平台。我们致力于成为业内领先的Triton技术交流与应用分享的殿堂,为推动人工智能技术的普及与深化应用贡献力量。
更多推荐
已为社区贡献1条内容
所有评论(0)