NeighborLoader on Hetrograph doesn't return the expected number of neighbors

[Sanjaykrishnamurthy]

Hi peers,
I noticed NeighborLoader on Hetrograph doesn't return the expected number of neighbors
Take the small example below, where i was expecting batch['item'].x to return the 1 neighbors of batch['user'].x but its empty. What am i doing wrong?

import torch
from torch_geometric.data import HeteroData
from torch_geometric.loader import NeighborLoader

# Number of nodes for each type
num_users = 4
num_items = 10

user_id = torch.arange(5)
item_id = torch.arange(10)

# Creating random edges (just for demonstration)
buy_edges = torch.randint(0, num_users, (2, 8), dtype=torch.long) 

# The second row of edges needs to be offset by num_users for item indices
buy_edges[1, :] += num_users\

data = HeteroData()

# Adding node data
data['user'].x = user_id
data['item'].x = item_id

# Adding edge data
data['user', 'buy', 'item'].edge_index = buy_edges

train_loader = NeighborLoader(
    data,
    num_neighbors={key: [1] for key in data.edge_types},
    batch_size=2,
    input_nodes=('user', data['item'].x),
)
batch = next(iter(train_loader))


print(batch['user'].x)
# tensor([0, 1])


print(batch['item'].x)
# tensor([], dtype=torch.int64)

[rusty1s]

The NeighborLoader samples from destination to source nodes (e.g., item to user in the first hop), and so you are missing a reverse direction:

data['item', 'bought_by', 'user'].edge_index = buy_edges.flip(0)

[Sanjaykrishnamurthy]

Thank you. Adding that line solved the issue.