Getting an error when trying to use dataloader for hetero link prediction with #4026

[Jeriousman]

Following #4026, I tried to make a dataloader for edge link prediction.
My train_data is as below:

HeteroData(
  card={ x=[15910, 300] },
  user={ x=[12363, 300] },
  (user, solved, card)={
    edge_index=[2, 66415],
    edge_label=[132830],
    edge_label_index=[2, 132830]
  },
  (card, rev_solved, user)={ edge_index=[2, 66415] }
)

Then I tried to make the dataloader as below:

train_loader = NeighborLoader(
    train_data,
    # Sample 30 neighbors for each node and edge type for 2 iterations
    num_neighbors={key: [5] * 2 for key in train_data.edge_types},
    # Use a batch size of 128 for sampling training nodes of type paper
    batch_size=128,
    input_nodes=('user__solved__card', train_data['user', 'solved', 'card'].edge_label_index.view(-1)))

I do not receive any error but when I do next(iter(train_loader)), it shows this as below,

sampled_hetero_data = next(iter(train_loader))

It shows this error.

Traceback (most recent call last):

  File "<ipython-input-159-cc85ad838780>", line 1, in <module>
    sampled_hetero_data = next(iter(train_loader))

  File "/home/hojun/anaconda3/envs/geometric/lib/python3.6/site-packages/torch_geometric/loader/base.py", line 21, in __next__
    return self.transform_fn(next(self.iterator))

  File "/home/hojun/anaconda3/envs/geometric/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 521, in __next__
    data = self._next_data()

  File "/home/hojun/anaconda3/envs/geometric/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 561, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration

  File "/home/hojun/anaconda3/envs/geometric/lib/python3.6/site-packages/torch/utils/data/_utils/fetch.py", line 52, in fetch
    return self.collate_fn(data)

  File "/home/hojun/anaconda3/envs/geometric/lib/python3.6/site-packages/torch_geometric/loader/neighbor_loader.py", line 89, in __call__
    self.directed,

IndexError: _Map_base::at

Am I doing something not according to #4026? Thank you in advance.

[rusty1s]

Oh you are right. This procedure really does not work yet for heterogeneous graphs, in particular because input_nodes currently expects a single node type as input. Sorry for the confusion. Will try to fix that in #4026 as well.

[Jeriousman]

Thank you in advance!