a = torch.tensor([[1,2,3,4],[5,6,7,8]])
idx = torch.tensor([[0,2,1],[2,3,0]])

# How to do it in batch ?
c_1 = a[0][idx[0]].view(1,-1)
c_2 = a[1][idx[1]].view(1,-1)
c = torch.cat((c_1, c_2), dim=0)

The desired output is:
tensor([[1, 3, 2],
        [7, 8, 5]])

I tried a[idx], however, it goes wrong.