Applying Embedding Layer to a specific node type

[IskaNeumann]

Hi,

I have a heterogeneous graph with 2 node types and a homogeneous model with two SAGEConv layers and a torch.nn.Embedding layer, which I convert to a heterogeneous model with .to_hetero(). I want to use the Embedding Layer to learn an embedded representation for a categorical feature of one Node Type. Currently I have the problem that in the forward pass of the model the embedding layer is called for both node types. This fails because the categorical feature is only available in one node type. Is it possible to apply the embedding layer only to a specific node type?

Thank you!

[rusty1s]

I think in this case it is best to out-source the embedding layer from the to_hetero call, e.g., via

class MyModel(torch.nn.Module):
    def __init__(self, ...):
         self.emb = Embedding(...)
         self.model = to_hetero(MyGNN(), metadata)
         
    def forward(self, x_dict, edge_index_dict):
         x_dict = copy.copy(x_dict)
         x_dict[key] = self.emb(x_dict[key])
         return self.model(x_dict, edge_index_dict)

[IskaNeumann]

thank you, that worked :)