Assume that I have a tensor A with dimension [Batch_Size, 5, 5]. I got an order tensor with dimension [Batch_size, 5]
Im trying to use the order to directly resort the tensor A like this:
A = A[Order, :]
However got the cuda error: out of memory. Is there any way to solve this?