[Torch.Sort Problem]

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?