How to use global_mean_pool for single graph dataset like Cora without batch?

[scottshufe]

Hi, I would like to use global_mean_pool function to get a single graph's embedding, but I got the following error information when I use global_mean_pool(x):
TypeError: global_mean_pool() missing 1 required positional argument: 'batch'

And if I set the param 'batch=None', I got another error:
AttributeError: 'NoneType' object has no attribute 'max'

So how should I use global_mean_pool correctly if I want to obtain the graph embedding given node embeddings from a single graph?

BTW, the version of my pytorch_geometric is 1.6.1.

Thanks. Look forward to your reply.

[rusty1s]

You can either do:

x = x.mean(dim=0, keepdim=True) if batch is None else global_mean_pool(x, batch)

or do

if batch is None:
    batch = torch.zeros(x.size(0), dtype=torch.long, device=x.device)
x = global_mean_pool(x, batch)

Note that this is already fixed in later PyG versions.

[scottshufe]

Thanks for your help!