Hi all,
when I am masking a sparse Tensor with index_select() in PyTorch 1.4, the computation is much slower on a GPU (31 seconds) than a CPU (~6 seconds).
Does anyone know why there is such a huge difference?
Here is a simplyfied code snippet for the GPU:
n= 2000
groups = torch.sparse_coo_tensor(indices= torch.stack((torch.arange(n), torch.arange(n)),
                                 values=torch.ones(n, dtype= torch.long),
                                 size=(n,n))
idx = torch.ones(1999, dtype= torch.long)

idx = idx.cuda()
groups = groups.cuda()

start_time = time.time()
groups.index_select(0, idx)
print("--- %s seconds for masking---" % (time.time() - start_time))