How to create mini-batches for weighted graphs given that LinkNeighborLoader doesn't take edge_weight?

[yuan-meng-dd]

I'm building a link prediction model for a weighted bipartite graph. Say the task is to predict if a query might be intended for a type of product and edge_weight captures the total clicks of each <query, product type> pair.

Since this is a link prediction task, I want to use LinkNeighborLoader to create mini-batches with a fixed number of links in each batch by doing something like this:

# load training data
train_loader = LinkNeighborLoader(
    data=train_data,
    num_neighbors=[20, 10],
    neg_sampling_ratio=2.0,
    edge_label_index=(("query", "intended_for", "taxonomy"), edge_label_index),
    edge_label=edge_label,
    
    batch_size=128,
    shuffle=True,
)

However, LinkNeighborLoader doesn't have edge_weight as an attribute. How else can I create mini-batches for a weighted graph? One solution is to repeated edges as many times as there are weights. However, search data is already in huge volume so repeating edges is not really feasible. Many thanks for any advice!