Suppose I have a tensor a of shape (2, 3, 4), and my order tensor idx is (2, 4). I need to rearrange it along the dim=2 dimension. Such that a[0] = a[0][:,idx[0]], a[1] = a[1][:,idx[1]].
My problem is that my tensor in reality is larger than the tensor in the example, and I would like to do it in single line which is faster and requires less memory.