How to use shuffle with train_mask in NeighborLoader()

[JihoChoi]

I am building a custom dataset that has two node types (heterogeneous setup). Type 1 nodes have labels (0,1, ..., |C|) and Type 2 nodes don't have labels (-1). Thus I set up a list of train_mask with the indexes of Type 1 nodes to be True and the other to be False. However when I use NeighborLoader with the shuffle flag = True, train_mask doesn't follow the shuffled index.

Could someone help me with this or help me find a good example of train_mask shuffle?

train_loader = NeighborLoader(
    dataset['train'],
    num_neighbors=[20] * 1,  # Sample 20 neighbors for each node and edge type for 2 iterations
    batch_size=100,
    # input_nodes=None,
    # shuffle=False,
    input_nodes=torch.tensor(dataset['train'].train_mask),
    shuffle=True,
)

Error (Due to the unshuffled index of train_mask)

return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
Target -1 is out of bounds.

[rusty1s]

You can just set the train_mask of a single node type:

data = HeteroData()
data['A'].train_mask = ...

loader = NeighborLoader(data, input_nodes=('A', data.['A'].train_mask), ...)