InnovArul:

array1d = torch.tensor([0, 5, 10, 20]) selection = array1d[index]


Thanks! This works much faster and is much cleaner than what I have right now. This was what I had before:
index = torch.tensor([[0,3],[2,1]])
array1d = torch.tensor([0, 5, 10, 20])
selection = array1d.reshape(-1).repeat(2,1).gather(1,index)
Yours is ~ 2.8X faster on my CPU