Obtaining the original edge indices from ClusterData

[realfolkcode]

Hello! I am currently struggling with link prediction problem on large graphs. I have decided to use ClusterGCN. Because my initial graph doesn't fit in memory, I compute partitioning using ClusterData and train embeddings on subgraphs. However, it is not clear how to infer after training. Is it possible to obtain the original edge indices of a subgraph after partitioning?

[rusty1s]

I suggest to add a vector to your data that denotes the index of original edges:

data.e_id = torch.arange(data.num_edges)

After applying ClusterLoader, this attribute will hold the original index of edges after partitioning.

[realfolkcode]

Nice! Thank you for answering! This should work well.