Index tensor with list

I have a tensor T of shape (x, y) and a list L of shape (x), containing numbers [0, y). I want to have a tensor of shape (x), where each the ith element is T[i, L[i]]. I need this to be differentiable, and ideally with only pure torch operations and without loops. Apparently, T[:, L] gives a tensor of shape (x, x), where every element of L indexes every row.

Currently, I’m using this horrendous workaround, which feels extremely inefficient:

R = (T[:, L] * torch.eye(x)).sum(dim=0)

This should work:

A, B = 10, 5
x = torch.randn(A, B)
idx = torch.randint(0, B, (A,))
x[torch.arange(A), idx]
6 Likes

Thanks, that worked!

@ptrblck
l=tensor([100, 12, 75, 55])
i have this tensor list

and i want to use it his way
X[:,- l:]
but i will get an error saying only
only integer tensors of a single element can be converted to an index

basically i want to mask certain range of labels/outputs while computing the loss/metric.
eg. for first tensor element in Batch i want to use elements starting only from -100:,for second element from -12: and like wise

I don’t think this would be possible via direct indexing as the result tensors would have a variable size, wouldn’t they?
In that case you could loop all indices and append the result tensors to a list.

Actually that would add to epoch time my bs is 1280 .
I would have loop over so many elements.

Not sure what can be other method of masking with variable no of elements to skip from every time