Solved by ptrblck in post #2 
              
              
                If I understand you correctly, you would like to select from the last dimension using an index tensor with values in [0, 2]. 
The same operation should be performed for all examples? 
If so, you could try the following code: 
a = torch.randn(20, 10, 3)

idx = torch.zeros(10).random_(3).long()
idx = …