NeighborLoader and edge weights

[NinaCalvi]

Hello!!
I am having a hard time understanding exaclty how the neighbor loader works
I am currently running something like this:

loader_train = NeighborLoader(
        dataset_train,
        num_neighbors = {('target', 'pred', 'node'): [0, 5],
                         ('node', 'pred', 'target'): [10, 0]}, 
        batch_size=1,
        input_nodes=input_nodes, drop_last=True,
        directed=True,
        replace=False,
        shuffle=False)

when I look at a produced batch I get the following:

> HeteroData(
>   node={ x=[10] },
>   target={
>     x=[2],
>     batch_size=1
>   },
>   (node, pred, target)={
>     edge_index=[2, 10],
>     edge_weight=[199]
>   },
>   (target, pred, node)={
>     edge_index=[2, 11],
>     edge_weight=[199]
>   }
> )

The edge indeces for each edge type are sampled in the way I'd expect them to, but the edge weights aren't aligned - the batch contains all the edge weights from the original graph. This is not what I want - I would be expecting 10 and 11 edge weights which correspond to the sampled indeces. Can someone help me here?

[NinaCalvi]

sorry solved!! the edge_indeces and edge_weight weren't of the same lenght initally - so this was causing issues!