PyTorch Geometric: How to use label propagation algorithm?

Hi there,

I am quite new to PyTorch (and geometric deep learning) and am trying to use the PyTorch geometric package to implement the ‘Label propagation’ algorithm for node classification.

Question: I am struggling to understand the documentation and don’t quite know how to implement it in code. Would someone be able to help me?

Context: The class is called LabelPropagation and is listed in the page.

Attempt: Should the ‘main’ portion of the implementation look like:

class LabelPropagation(torch.nn.Module):
def __init__(self, num_layers, alpha):
        super(LabelPropagation, self).__init__()
        self.num_layers = num_layers
        self.alpha = alpha

    def forward(self, x, edge_index):
        # I don't really know what to put here
        return x

Any help would be greatly appreciated as I am still learning how to understand and interpret documentation.