data = Data(edge_index=edge_index)
data.num_nodes = edge_index.max().item() + 1
data.x = torch.eye(data.num_nodes, dtype=torch.float)
y = [0 if G.nodes[i]["club"] == "Mr. Hi" else 1 for i in G.nodes]
data.y = torch.tensor(y)
self.data, self.slices = self.collate([data])
After Change
// Select a single training node for each community
// (we just use the first one).
train_mask = torch.zeros(y.size(0), dtype=torch.bool)
for i in range(int(y.max()) + 1):
train_mask[(y == i).nonzero(as_tuple=False)[0]] = True
data = Data(x=x, edge_index=edge_index, y=y, train_mask=train_mask)