Python `寄存器_forward _hook `在使用多个GPU时出错

我正在使用

register\u forward\u hook

从另一个模型中提取中间特征，并在另一个模型中使用此特征。然而，该代码在单个GPU上运行良好，但出现了一个

运行时错误：预期'out'的张量与参数#2'mat1'的张量具有相同的设备；但是设备3不等于4（在检查addmm的参数时）

这是我的密码。谁能告诉我怎么解决这个问题吗

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

device_id = [3,4]
data_size = 2000
hidden_length = 10
batch_size = 100
device_ = 3
input_size = 1000
output_size = 200
device = torch.device("cuda:{}".format(device_) if torch.cuda.is_available() else "cpu")

class RandomDataset(Dataset):

    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),
                         batch_size=batch_size, shuffle=True)


class Model(nn.Module):
    # Model to be extracted

    def __init__(self, input_size, output_size, hidden_length):
        super(Model, self).__init__()
        block = [nn.Linear(input_size, input_size)] * hidden_length
        self.hidden = nn.Sequential(*block)
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.hidden(input)
        output = self.fc(output)

        return output


class HookModel(nn.Module):
    # Our model

    def __init__(self, input_size, output_size, hidden_length):
        super(HookModel, self).__init__()
        self.model = Model(input_size, output_size, hidden_length)
        self.layer_extract = self.model._modules.get("hidden")
        self.fc = nn.Linear(input_size, input_size)

    def forward(self, input):
        def model_forward(input):
            model_output = []

            def hook(m, i, o):
                model_output.append(o.clone().detach())

            h = self.layer_extract.register_forward_hook(hook)
            self.model(input)
            h.remove()
            return model_output[0]

        output = model_forward(input)
        output = self.fc(output)

        return output


model = HookModel(input_size, output_size, hidden_length)
if len(device_id) > 1:
    print("Let's use", len(device_id), "GPUs!")
    model = nn.DataParallel(model, device_ids=device_id)
model.to(device)

for data in rand_loader:
    input = data.to(device)
    output = model(input)
    loss = torch.sum((output - 2) ** 2)
    loss.backward()