Solved by ptrblck in post #2 
              
              
                You could directly index the tensor and permute the output: 
res = a[torch.arange(2), :, index]
print(res)
> tensor([[1, 3],
          [4, 6]])
print(res.t())
> tensor([[1, 4],
          [3, 6]])