class GCN(torch.nn.Module): def __init__(self): super().__init__() self.gcn = GCNConv(dataset.num_features, 3) self.out = Linear(3, dataset.num_classes)