Efficient way to select and collapse tensor along specific axis

1

Hi.

I am finding an efficient way to select and collapse tensor along specific axis.

Suppose and I have a Tensor X whose shape is (N, L, D) and a list of ids I whose length is N.
What I need is a Tensor Y of shape (N, D) where Y[n, :] = X[n, I[n], :] .

I found torch.gather can do this kind of thing like this.

idx = torch.LongTensor(N, 1, D)
for n in range(N):
    idx[n, :, :] = I[n]

Y = torch.gather(X, 1, idx)

However this code is obviously not efficient.
It creates a new tensor and fill it by for loop.

Are there any better way?

Thank you.