pytorch采坑（Expected all tensors to be on the same device, but found at least two devices）

问题描述

pytorch版本：2.0.1+cu117

使用pytorch的时候，我将模型部署到gpu上，然后将输入也部署到gpu上，但是却出现了如下报错

贴上我的源代码：

import torch.nn as nn
import torch.nn.functional as F
import torch


class Modeltest(nn.Module):

    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv1d(1, 20, 5)
        self.conv2 = nn.Conv1d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

input = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9,
                      10]).to(device).type(torch.FloatTensor)

input = torch.reshape(input, (1, -1))

modeltest = Modeltest().to(device)

output = modeltest(input)

output2 = output.sum() * 0.1

print(output.device)
print(output2.device)

从代码中可以看出，我定义了input，并写了to(device)，然后对input进行了reshape操作，结果出现了报错

解决方案

我们最开始定义的input确实是在gpu上，但是使用了reshape之后，相当于重新定义了一个input，而默认情况下，新定义的变量都是在cpu上的，所以出现了问题，正确的代码应该如下：

import torch.nn as nn
import torch.nn.functional as F
import torch


class Modeltest(nn.Module):

    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv1d(1, 20, 5)
        self.conv2 = nn.Conv1d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

input = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9,
                      10]).type(torch.FloatTensor)

input = torch.reshape(input, (1, -1)).to(device)

modeltest = Modeltest().to(device)

output = modeltest(input)

output2 = output.sum() * 0.1

print(output.device)
print(output2.device)

或者第一个input定义的时候直接使用reshape，也可以，如下：

import torch.nn as nn
import torch.nn.functional as F
import torch


class Modeltest(nn.Module):

    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv1d(1, 20, 5)
        self.conv2 = nn.Conv1d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

input = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9,
                      10]).type(torch.FloatTensor).reshape(1, -1).to(device)

modeltest = Modeltest().to(device)

output = modeltest(input)

output2 = output.sum() * 0.1

print(output.device)
print(output2.device)