If I correctly understood your question, this would work. If you want a more general way, you can easily extend it with a loop. It was not a matter of PyTorch, just Python programming…
import torch
a = torch.tensor([[1,2,3],[4,5,6],[7,8,9]])
b = torch.tensor([1,2,1])
rearranged_a = torch.zeros(a.size())
rearranged_a[:2] = a[b==1]
rearranged_a[-1] = a[b==2]