ShaDowKHopSampler creates incorrect data.batch when using Batch graph

🐛 Describe the bug

Hello! Firstly, thank you for this fantastic library.

I am trying to run various scalable GNNs on custom data but I am having an issue with shaDow.

On this line in the shaDow model definition (as per the example): x = torch.cat([x[root_n_id], global_mean_pool(x, batch)], dim=-1) This error is thrown: RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 1024 but got size 1 for tensor number 1 in the list. Where 1024 is the batch size.

I’ve found this is only a problem when the graph dataset is composed into a disjointed Batch graph.

To reproduce the problem, run the code from examples/shadow.py but change the data object into a Batch object:

from torch_geometric.data import Batch
dataset = Flickr(path)
datas = [dataset[0]]
data = Batch.from_data_list(datas)

Thank you for any help!

Environment

• PyG version: 2.1.0.post1
• PyTorch version: 1.11.0+cu113
• OS: Ubuntu 20.04.2 (WSL) and macOS 11.1
• Python version: 3.7.2
• CUDA/cuDNN version: 11.3
• How you installed PyTorch and PyG (conda, pip, source): pip
• Any other relevant information (e.g., version of torch-scatter):

[EdisonLeeeee]

Thank you for making it clear. I get your point now. Previously, data.num_graphs is calculated by batch.max() + 1, which would lead to a wrong result if batch was overridden in the sampler (e.g., passing a Barch object).

Now I think data.num_graphs would return a correct result even passing a Batch object after the fix in PR. I’ve tested the example/shadow.py example with your case, and it works fine – data.num_graphs always equals out.size(0).

[cvanea]

Thanks very much! I’m closing the Issue then 😃