How can I assign the edge-specific attribute from the LinkNeighborLoader?

[songsong0425]

Hi, always thank you for your help.
I have a question about the approaches for the edge class labeling.
When I have 2 types of edge classes and want to allocate the related class ID to the sampled neighbors from LinkNeighborLoader, how can I do it?
I want to do it to separate the message passing per edge classes.

Please take a look at this example. (the source: here)

For example, I already have the edge class information about the answer sets (red and blue lines in the image), and the neighbors were sampled using the LinkNeighborLoader.
In this case, I aim to assign the edge class ID to the neighbors based on the seed edges.
The data would look like this:

# Example form of Data object for the left image
# I guess that batch size is 2 since the number of edges with edge labels is 2
# The original Data result (Note that, the network in the image is undirected, but I consider it as the directed network for simplicity)
Data(x=[10, 128], edge_index=[2, 8], edge_label=[2], edge_label_index=[2,2], edge_class=[2], num_nodex=10, n_id=[10], e_id=[8], num_sampled_nodes=[2], num_sampled_edges=[1], input_id=[2])

# The expected Data result
Data(x=[10, 128], edge_index=[2, 8], edge_label=[2], edge_label_index=[2,2], edge_class=[8], num_nodex=10, n_id=[10], e_id=[8], num_sampled_nodes=[2], num_sampled_edges=[1], input_id=[2])

print(edge_index)
tensor([[1, 3, 4, 2, 6, 8, 9, 7],
        [0, 0, 0, 1, 5, 5, 5, 6]])
print(edge_class)
tensor([1, 1, 1, 1, 2, 2, 2, 2])

Before I tried using input_id for the same task, but it only applied to the edges having edge_label.

data.edge_class = data.edge_class[data.input_id]

Also, I looked subgraph in the PyG, but I'm not sure how to use it for this job.

I'm looking forward to your suggestion.
Thank you for reading this question.

---

Additional information: I used the LinkNeighborLoader as below:

train_loader = LinkNeighborLoader(kf_train_data, edge_label_index=kf_train_data.edge_label_index, edge_label=kf_train_data.edge_label,
                                      batch_size=nbatch, shuffle=True, neg_sampling_ratio=0.0, num_neighbors=[-1, -1]) 
...