To be honest, I’ve never used the torch_geometric Data class so I don’t know for sure. To get a subset of indices from a tensor T, you’d do the following:

T = torch.randn(10,5,5)                                    # create some dummy data for example
idx = torch.tensor([1,2,4,7],dtype=torch.uint8) # make tensor of desired indices
T.shape                                                            # returns (10,5,5)
T[idx].shape                                                     # returns (4,5,5)

I’d guess that the way the Data class is defined, grabbing a subset of the indices in a Data object will slice each of the attributes in Data in the way defined above, though I can’t be 100% sure and I’m too lazy to download the packages to check.
You should try using the format above, where idx is a 1D tensor with the indices of your desired train_mask and test_mask.