How can I efficiently implement this code? Implementing without the loop.

import torch
torch.manual_seed(0)
tensor = torch.rand(2,4,3,3)
def sort_tensor(tensor):
        # average pooling
        polling = torch.nn.AdaptiveAvgPool2d(1)
        tensor_pooled = polling(tensor)
        # sorting the averaged pooled values and getting the indexes 
        sorted_tensor_pooled, sorted_indx = torch.sort(tensor_pooled, 1, descending=True)
        sorted_indx = sorted_indx.squeeze()
        
        #sorting the orginal tensor based on the sroting index
        for i in range(len(sorted_indx)):
            if i == 0:
                sorted_tensor = tensor[i, sorted_indx[i], :, :].unsqueeze(dim=0)
            else:
                sorted_tensor = torch.cat((sorted_tensor, tensor[i, sorted_indx[i], :, :].unsqueeze(dim=0)), dim=0)

        return sorted_tensor
print('sorted', sort_tensor(tensor))