Batch tensor not being put on cuda

[davidireland3]

I am using top_k pooling with GPU but having an issue where the batch tensor which tells which graph the nodes belong to, from the data loader, is not being put on the GPU.

The code looks like

    def forward(self, data, return_node_embedding=False, return_both=False):
        x, edge_index, batch = data.x, data.edge_index, data.batch  # here data.batch is not on the gpu like the x and edge index tensors

        x = torch.relu(self.gcn1(x, edge_index))
        if return_node_embedding and not return_both:
            return x
        elif return_both:
            x_ = copy.deepcopy(x.detach())
        x, edge_index, _, batch, _, _ = self.first_pool_layer(x, edge_index, batch=batch) 

The data comes from a DataLoader using train_loader = DataLoader(train_data, batch_size=128, shuffle=True).

Has anyone had this issue?

e: this has happened since updating to the latest version of pyg.

[rusty1s]

This should be already fixed in master, see #3214.

[davidireland3]

as far as I know I was: I used the pip command that you shared in my conda environment so it should've been the master.

also, in case you are interested, I was still having some problems with the embeddings I mentioned in an earlier comment but I narrowed the problem down to being on PyTorch version 1.9 (as opposed to 1.8) which I also updated at the same time as the new PyG version -- not sure why this caused the problem but at least I found it.

[rusty1s]

Can you still confirm that the above code snippet from me works as expected for you?

[davidireland3]

I can confirm it doesn't work as intended. My output was:

cuda:0
cpu
cuda:0
cpu
cuda:0
cpu
cuda:0
cpu
cuda:0
cpu
cuda:0
cpu

[rusty1s]

I think your old conda installation of PyG is still the active one. You should try to uninstall it first (conda uninstall pyg), and re-install from pip again.

[davidireland3]

Ah yes, this fixed it! Thanks so much for the help :-)