import torch
a = torch.tensor([[1,2,3,4],[5,6,7,8]])
idx = torch.tensor([[0,2,1],[2,3,0]])
idx2 = idx + torch.arange(idx.size(0)).view(-1, 1) * a.size(1)
c = a.view(-1)[idx2]

It works as long as a and idx have only two dimensions and idx.size(0) <= a.size(0).