Random edge sampling

[GianlucaDeStefano]

Hello, I am trying to create a Link prediction model and have a question about how to sample training edges properly.

Essentially, in my problem, each node has an embedding extracted by a natural language model and using it (and only it) I want to predict whether an edge should be present between two nodes.

I am not interested in information from connected Neighbours, as I will not use any graph-based convolution.
All I want is to sample random edges, get the two connected nodes and use their local features to predict if the link should be present.

Is there any Dataloader that will allow me to do this? At the moment I am using LinkNeighborLoader in this way:
LinkNeighborLoader(dataset_train.data, num_neighbors = [0] , shuffle = True ,batch_size = batch_size)

The issue is that the batch.edge_index tensor I get in this way is empty.
Indeed by using a batch size of 16 I obtain the following batch:
Data(x=[32, 2, 512], edge_index=[2, 0], input_id=[16], edge_label_index=[2, 16])

Note: edge_index=[2, 0]

Am I doing something wrong?
Is there any Dataloader Class more suited to my use case?

Thank you in advance for any :).

[wsad1]

No you aren't doing anything wrong. edge_index is used to run message passing, since you specified num_neighbors=[0] the loader simply doesn't sample any edges for message passing. edge_label_index can be used for the link prediction task.

Since you don't want neighborhood information its probably better to use torch.DataLoader.

[GianlucaDeStefano]

Ahhh I see, I miss-interpreted the meaning of edge_label_index . Thanks!