Solved by ptrblck in post #2 
              
              
                This should work: 
A = torch.tensor([[[1,2],[3,4],[5,6],[7,8]] , [[9,10],[11,12],[13,14],[15,16]]])
index = torch.tensor([0,2])

out = A[torch.arange(A.size(0)), index]
print(out)
> tensor([[ 1,  2],
          [13, 14]])