How to return the intermediate result from propagate() function?

[Elnath-123]

I wrote a graph model class (Inherited from Message Passing class), and I rewrote the message() function to calculate the information needed to be passing from neighbor nodes to target node. For some reasons in my experiments, I need to get the intermediate results calculated in the message() function. I checked some tutorial codes, in which the propagate() function only return a variable named "out" (representing the encoded node feature). How can I get the intermediate results(except the variable "out") from propagate() function? I want to get the output of the last hidden layer of self.mlp_1.

def __init__(self, node_dim, out_channels, edge_dim):
    super(GraphModel, self).__init__()
    self.negative_slope = 0.2
    self.lin_node = nn.Linear(node_dim, out_channels)
    self.lin_edge = nn.Linear(edge_dim, out_channels)
    self.mlp_1 = nn.ModuleList([
        nn.Linear(out_channels, out_channels),
        nn.ReLU(),
        nn.Linear(out_channels, out_channels),
        nn.ReLU(),
        nn.Linear(out_channels, 1)
    ])
    self.att = Parameter(torch.Tensor(1, out_channels))
def forward(self, x, edge_index, edge_attr):
    x = self.lin_node(x)
    out = self.propagate(edge_index, x=x, edge_attr=edge_attr)
    return out

def message(self, x_i, x_j, edge_attr, index, ptr):
    edge_attr = self.lin_edge(edge_attr)
    fea = x_i + x_j + edge_attr
    hidden_fea = []
    length = len(self.mlp_1)
    for i in range(length):
        fea = self.mlp_1[i](fea)
        if i == length - 3:
            hidden_fea.append(fea)
    hidden_fea = torch.cat(hidden_fea, dim=1)

    fea = F.leaky_relu(fea, self.negative_slope)

    alpha = (fea * self.att).sum(dim=-1)
    alpha = softmax(alpha, index, ptr)
    self._alpha = alpha
    return (x_j + hidden_fea) * alpha.unsqueeze(-1)

[wsad1]

You could save the intermediate results as an instance variable self. and access it in the forward method after calling propogate. For an example check out the alpha variable in TranformerConv, its an intermediate result in the message function.