Graclus doesn't seem to be working with GPU

[mdanb]

I'm trying to use Graclus along with my GPU but there seems to be an issue (which doesn't happen if I use the CPU):

            for layer in range(1, self.num_layers):
                h1 = self.norms[layer - 1](h)
                h2 = F.relu(h1)
                h2 = F.dropout(h2, p=self.dropout, training=self.training)
                h = self.gcns[layer](h2, edge_index) + h
                g = max_pool(graclus(edge_index), Data(x=h, edge_index=edge_index, batch=batch))
                h = g.x
                edge_index = g.edge_index
                batch = g.batch

but I get

    g = max_pool(graclus(edge_index), Data(x=h, edge_index=edge_index, batch=batch))
  File "/home/mdanb/.local/lib/python3.8/site-packages/torch_geometric/nn/pool/graclus.py", line 35, in graclus
    return graclus_cluster(edge_index[0], edge_index[1], weight, num_nodes)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "/home/mdanb/.local/lib/python3.8/site-packages/torch_cluster/graclus.py", line 44, in graclus_cluster
    # Randomly shuffle nodes.
    if weight is None:
        perm = torch.randperm(row.size(0), dtype=torch.long, device=row.device)
               ~~~~~~~~~~~~~~ <--- HERE
        row, col = row[perm], col[perm]
RuntimeError: radix_sort: failed on 1st step: cudaErrorInvalidDevice: invalid device ordinal

PS: This doesn't seem to be an issue for CUDA versions < 11.1
It worked with 10.1

[rusty1s]

I'm not exactly sure - This might be a PyTorch issue. Does the same error happen if your remove the @torch.jit.script annotation in .../torch_cluster/graclus.py?

[mdanb]

Yup the error is still there

[rusty1s]

That is weird. This should at least be a different error (not related to the TorchScript interpreter).

[mdanb]

Wait it is actually somewhat different sorry:

Traceback (most recent call last):
  File "/home/mdanb/deep_gcns_torch/examples/ogb/ogbg_ppa/main.py", line 209, in <module>
    main()
  File "/home/mdanb/deep_gcns_torch/examples/ogb/ogbg_ppa/main.py", line 158, in main
    epoch_loss = train(model, device, train_loader, optimizer, criterion)
  File "/home/mdanb/deep_gcns_torch/examples/ogb/ogbg_ppa/main.py", line 34, in train
    pred = model(batch)
  File "/home/mdanb/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/mdanb/deep_gcns_torch/examples/ogb/ogbg_ppa/model.py", line 104, in forward
    g = max_pool(graclus(edge_index), Data(x=h, edge_index=edge_index, batch=batch))
  File "/home/mdanb/.local/lib/python3.8/site-packages/torch_geometric/nn/pool/graclus.py", line 35, in graclus
    return graclus_cluster(edge_index[0], edge_index[1], weight, num_nodes)
  File "/home/mdanb/.local/lib/python3.8/site-packages/torch_cluster/graclus.py", line 44, in graclus_cluster
    perm = torch.randperm(row.size(0), dtype=torch.long, device=row.device)
RuntimeError: radix_sort: failed on 1st step: cudaErrorInvalidDevice: invalid device ordinal