HeteroSAGEConv - Issue with edge_index in message_passing when selecting indices from data

I have a graph with two different types of nodes (user and doc). The edge_index contains edges between users and also users and docs. In message_passing, node_feature_neigh is getting assigned to data (data = kwargs.get(arg[:-2], inspect.Parameter.empty)) when arg is 'node_feature_neigh_j'. However, later, index_select is used to select rows from data using node_id in edge_index which can be either a user or a doc. So the issue is that data at this point contains only the features for doc nodes but edge_index contains node_ids from users as well and therefore, I get index out of range exception.

  File "...torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "...deepsnap/deepsnap/hetero_gnn.py", line 135, in forward
    self.convs[message_key](
  File "...torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "...deepsnap/deepsnap/hetero_gnn.py", line 44, in forward
    return self.propagate(
  File "...torch_geometric/nn/conv/message_passing.py", line 250, in propagate
    kwargs = self.__collect__(edge_index, size, mp_type, kwargs)
  File "...torch_geometric/nn/conv/message_passing.py", line 147, in __collect__
    out[arg] = data.index_select(self.node_dim,
IndexError: index out of range in self

In the provided example for heterogenous graph, cora and citseer are combined in one graph where two graphs are isolated and each have their own node and edge types and I guess that's why this code works.

I would really appreciate any help with this. Happy to elaborate more if necessary.

snap-stanford / deepsnap

HeteroSAGEConv - Issue with edge_index in message_passing when selecting indices from data #8