You misunderstood me: I just meant that you should use a for loop but only if you have to do more than two replacements and just for compactness. I also would say that how vector b is built is rather counterintuitive but I don’t know from where it comes. I can guess some class labels. However, if instead of having b like this you manage to have b = [0,2,1], which gives the indeces where you want to place the rows, there exists a builtin method that does the job. Try, with b = [0,2,1]:
rearranged_a = torch.index_select(a,0,b)