PyTorch嵌套nn.Module迁移GPU后Embedding层报CPU运算错误如何解决

你所有自定义nn.Module子类的__init__方法均未调用父类nn.Module的构造函数，导致子模块、可训练参数没有被正确注册到PyTorch的模块管理体系中，因此调用.to(device)时无法自动将所有子层参数迁移到指定设备。

• 给每个自定义模块的__init__方法第一行补充父类初始化调用super().__init__()，修正后的代码示例如下：

import torch.nn as nn

class M_outer(nn.Module):
    def __init__(self, **kwargs):
        super().__init__()
        self.inner_1 = M_inner(**kwargs)
        self.inner_2 = M_inner(**kwargs)
        # ...
    
    def forward(self, input):
        result = self.inner_1(input)
        result = self.inner_2(result)
        # ...
        return result
    

class M_inner(nn.Module):
    def __init__(self, **kwargs):
        super().__init__()
        self.sub_1 = M_sub(**kwargs)
        self.sub_2 = M_sub(**kwargs)
        # ...
    
    def forward(self, input):
        result = self.sub_1(input)
        result = self.sub_2(result)
        # ...
        return result

class M_sub(nn.Module):
    def __init__(self, x, y,** kwargs):
        super().__init__()
        self.emb_1 = nn.Embedding(x, y)
        self.emb_2 = nn.Embedding(x, y)
        # ...
        self.norm  = nn.LayerNorm(y)
    
    def forward(self, input):
        emb = (self.emb_1(input) + self.emb_2(input))        
        # ...
        return self.norm(emb)

• 确认输入数据张量也迁移到了对应设备，模型迁移到GPU后，输入张量也需要调用.to(device)之后再传入模型推理/训练，否则也会出现跨设备运算报错：

input = input.to(device)
output = model(input)

• 可以通过以下代码验证所有参数的设备，确认是否全部迁移成功：

for name, param in model.named_parameters():
    print(f"参数名：{name}，所在设备：{param.device}")

nn.Embedding本身就是nn.Module的子类，不需要你手动用nn.Parameter注册，只要你的自定义模块正确调用了父类构造函数，直接赋值为实例属性的nn.Embedding会被自动识别为子模块，其内部参数会被纳入PyTorch的统一管理，调用.to()、.train()、.eval()等方法时都会同步生效。

内容的提问来源于stack exchange，提问作者SumNeuron