Retrieving individual Node embeddings of a batch of graphs

[DogitoErgoSum]

Hi!

In my task i need to mix node and graph classification. I want to predict node features of separate graphs.
Therefore, i want to apply message propagation in a batch of graphs and retrieve the node embedding of each node but i need to map where each node belongs.

I think i know how to do this but i am not sure if its the true behaviour

Take this example:

class GCN(torch.nn.Module):
    def __init__(self, hidden_channels):
        super(GCN, self).__init__()
        torch.manual_seed(12345)
        self.conv1 = GCNConv(output_size, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
    
        self.lin = Linear(hidden_channels, output_size)

    def forward(self, x, edge_index, batch):
        # 1. Obtain node embeddings 
        x = self.conv1(x, edge_index)
        x = x.relu()
        x = self.conv2(x, edge_index)
        

        # 2. Readout layer
        
        x = self.lin(x)
        
        return x

model = GCN(hidden_channels=64)

My batch data has 5 graphs :

If i pass my batch data to the model :

Does the first row of the output tensor belong to the first graph? And does the 4th row of my output tensor belong to the third node of my second graph?

Thank you

[Padarn]

Does the first row of the output tensor belong to the first graph? And does the 4th row of my output tensor belong to the third node of my second graph?

Correct. The batch is the 'graph index'.