1

我是 PyTorch 的新手，有一些nn.Module我想在 GPU 上运行的自定义。我们称它们M_outer为 、M_inner和M_sub。通常，结构如下所示：

class M_outer(nn.Module):
    def __init__(self, **kwargs):
        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):
        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, **kwargs):
        self.emb_1 = nn.Embedding(x, y)
        self.emb_2 = nn.Embedding(x, y)
        # ...
        self.norm  = nn.LayerNorm()
    
    def forward(self, input):
        emb = (self.emb_1(input) + self.emb_2(input))        
        # ...
        return self.norm(emb)

我尝试通过以下方式将我的模块放在 gpu 上：

model = M_outer(params).to(device)

然而，我仍然从嵌入层收到错误，说某些操作在 cpu 上。

我已阅读文档。我已经阅读了像这样有用的讨论帖子和像这样的相关 StackOverflow 帖子。

我无法注册nn.EmbeddingLayervia nn.Parameter。我错过了什么？

4

0 回答 0