Torch_geometric - Calculate Dirichlet energy at each layer

Hi, I’m pretty new to GNNs and PyTorch Geometric. I want to calculate the Dirichlet energy at each layer and plot them, so I can visualize the over-smoothing problem. I have looked into several ways to accomplish this, but none of them seem to work, as I have to compute the product between the features and graph Laplacian. I am not sure how to get the Laplacian when the hidden layer changes the dimension of the node features, as the dimensions of the new features and Laplacian computed off the edge_index would be completely different. Any tips pointing me in the right direction would be helpful!

I have a pretty standard GCN defined here:

class GCN(nn.Module):
    def __init__(
        self,
        n_hidden,
        n_features,
        n_classes,
        hidden_dim,
        dropout_rate,
        activation=F.relu,
        use_linear=False,
    ):
        super().__init__()
        self.n_hidden = n_hidden
        self.hidden_dim = hidden_dim
        self.activation = activation
        self.use_linear = use_linear

        self.in_conv = GCNConv(n_features, self.hidden_dim)
        self.hidden_convs = nn.ModuleList(
            [
                GCNConv(self.hidden_dim, self.hidden_dim)
                for _ in range(self.n_hidden)
            ]
        )
        self.out_conv = GCNConv(self.hidden_dim, n_classes)
        self.dropout = Dropout(dropout_rate)

    def forward(self, x, edge_index):
        x = self.dropout(x)
        x = self.activation(self.in_conv(x, edge_index))

        for i in range(self.n_hidden):
            x = self.activation(self.hidden_convs[i](x, edge_index))
            if self.use_linear:
                linear = Linear(self.hidden_dim, self.hidden_dim)(x)
                x = x + linear

        x = self.activation(self.out_conv(x, edge_index))
        x = self.dropout(x)

I would also like to accomplish this with a GAT and GIN as well.