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 = a[:,idx], a = a[:,idx].
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.
How could I do this please ?