Should I use cuda tensors for indexing if the data is on cuda?

1

Suppose an nn.module model has some indexing work happen in forward():

def forward(self, data):
    # some indexing works, e.g.,:
    index = torch.randint(4, (4,))  # should I add index = index.cuda()?
    self.linear(data[:, index])
    ...

If this model is on GPU, for the best performance, should I also move index to cuda? Or it doesn’t matter? Why?

1 Like
2

If data is on the GPU, you have to move the indices to the GPU. Otherwise, it will raise an error.

3

As an example shown below, I don’t move index to the GPU but the code still runs. My question is: Will adding index = index.cuda() help improve performance? and why?

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.linear = nn.Linear(1000, 200)
        
    def forward(self, data):
        index = torch.randint(1000, (1000,))  # should I add index = index.cuda()?
        return self.linear(data[:, index])
    
a = MyModule().cuda()
b = torch.randn(1000, 1000).cuda()
a(b)
4

Ho,

I guess in this case, it considers the Tensor as just a sequence of numbers.
Moving the indices to CUDA will definitely be faster.

5

Thanks. You are right.