Selecting over dimension by indices

1

Given a tensor of size 4x5x100 (rows x cols x k) I want select a different column for each rows resulting in a 4x100 tensor.

For example for row 0 I want to chose column 2, for row 1 I want to chose column 4, for row 2 I want to chose column 0 and so on!

Is there a single operation in Pytorch which takes a tensor and the indices I want to chose for dimension 1?

2

A bit ugly, but for now this should work:

a = torch.randn(4, 5, 100)
indices = torch.LongTensor([2, 4, 0, 1])
torch.gather(a, 1, indices.view(-1, 1).unsqueeze(2).repeat(1, 1, 100))
7 Likes
3

Thank you! At least it works :slight_smile: