The question about function Topk

[jiaruHithub]

When I use a function for experiments, I want to select top K scores:

import torch
from torch_geometric.nn.pool.topk_pool import topk
from torch.nn import Parameter
from torch_scatter import scatter_add, scatter_max

score = torch.rand(18)#.reshape(10, -1)
batch = torch.range(0,2).repeat_interleave(6,0).type(torch.int)

perm = topk(score, 0.1, batch)

But the error is reported:

Traceback (most recent call last):
  File "trails.py", line 72, in <module>
    perm = topk(score, 0.1, batch)
  File "/opt/conda/lib/python3.7/site-packages/torch_geometric/nn/pool/topk_pool.py", line 18, in topk
    num_nodes = scatter_add(batch.new_ones(x.size(0)), batch, dim=0)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "/opt/conda/lib/python3.7/site-packages/torch_scatter/scatter.py", line 22, in scatter_add
            size[dim] = int(index.max()) + 1
        out = torch.zeros(size, dtype=src.dtype, device=src.device)
        return out.scatter_add_(dim, index, src)
               ~~~~~~~~~~~~~~~~ <--- HERE
    else:
        return out.scatter_add_(dim, index, src)
RuntimeError: index 4294967297 is out of bounds for dimension 0 with size 3

Can you tell me why?Thanks!

[wsad1]

With batch as type(torch.long) the code you shared runs fine for me Which version of torch_scatter and pyg are you using?

[jiaruHithub]

Thanks a lot! It is fixed, my version:

torch                  1.5.0
torch-cluster          1.5.5
torch-geometric        1.5.0
torch-scatter          2.0.5
torch-sparse           0.6.6
torch-spline-conv      1.2.0

Why does this error is reported? Thanks!