Link Classification with LinkNeighborLoader and edge_attr Misalignment

[fafal-abnir]

I'm working on link classification using a Graph Neural Network with LinkNeighborLoader.
In my model, I classify edges by computing:

h_src = h[edge_index[0]]
h_dst = h[edge_index[1]]
combined = torch.cat([h_src, h_dst, edge_attr], dim=1)
out = self.out_mlp(combined)

The LinkNeighborLoader provides edge_attr for the sampled edge_index,
but not for the edge_label_index, which contains the edges I actually need to classify.
for example:

Data(
    x=[890, 16],
    edge_index=[2, 2219],
    edge_attr=[2219, 8],
    n_id=[890],
    e_id=[2219],
    num_sampled_nodes=[3],
    num_sampled_edges=[2],
    input_id=[64],
    edge_label_index=[2, 64],
    edge_label=[64]
)

• edge_attr → shape [2219, 8] — corresponds to edge_index
• edge_label_index → shape [2, 64] — edges I need to classify
• Missing: There is no edge_attr for edge_label_index in the batch.

Currently, I have to look up edge_attr from the full graph using edge_label_index to get the corresponding edge features.
This is inefficient and breaks the purpose of minibatch training.
Is there a better solution to this?

Specifically:

How can I efficiently obtain or include `edge_attr` for edge_label_index in each mini-batch?

Is there a built-in way in `LinkNeighborLoader` to support this use case?

Or is there an alternative strategy recommended for link classification with edge features?