I am analyzing example from pytorch-geometric library and I have a question.
How does graph network transform MNIST data into probability vector in pytorch geometric?
Here is the GCN:
def normalized_cut_2d(edge_index, pos):

    row, col = edge_index

    edge_attr = torch.norm(pos[row] - pos[col], p=2, dim=1)

    return normalized_cut(edge_index, edge_attr, num_nodes=pos.size(0))

class GCN1(torch.nn.Module):

    def __init__(self, d):

        print(d.num_features)

        super(GCN1, self).__init__()

        self.conv1 = SplineConv(d.num_features, 32, dim=2, kernel_size=5)

        self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5)

        self.fc1 = torch.nn.Linear(64, 128)

        self.fc2 = torch.nn.Linear(128, d.num_classes)

    def forward(self, data):

        data.x = F.elu(self.conv1(data.x, data.edge_index, data.edge_attr))

        weight = normalized_cut_2d(data.edge_index, data.pos)

        cluster = graclus(data.edge_index, weight, data.x.size(0))

        data.edge_attr = None

        data = max_pool(cluster, data, transform=transform)

        data.x = F.elu(self.conv2(data.x, data.edge_index, data.edge_attr))

        weight = normalized_cut_2d(data.edge_index, data.pos)

        cluster = graclus(data.edge_index, weight, data.x.size(0))

        x, batch = max_pool_x(cluster, data.x, data.batch)

        x = global_mean_pool(x, batch)

        x = F.elu(self.fc1(x))

        x = F.dropout(x, training=self.training)

        return F.log_softmax(self.fc2(x), dim=1)