Indexing on axis by list in PyTorch
I have Variables lengths_X of size (10L,) and A of size (10L, 16L, 5L).

I want to use lengths_X to index along the second axis of A. In other words, I want to get a new tensor predicted_Y of size (10L, 5L) that indexes axis 1 at i for all entries with index i in axis 0.

What is the best way to do this in PyTorch?

What you are looking for is actually called batched_index_select and I looked for such functionality before but couldn't find any native function in PyTorch that can do the job. But we can simply use:

A = torch.randn(10, 16, 5)
index = torch.from_numpy(numpy.random.randint(0, 16, size=10))
B = torch.stack([a[i] for a, i in zip(A, index)])


You can see the discussion here. You can also check out the function batched_index_select provided in the AllenNLP library. I would be happy to know if there is a better solution.
