Hi @Abe , I have converted the loop into a vectorized format.
I have verified with multiple inputs, and it works much faster.
Here’s the code.
def sort_tensorV2(tensor):
        # average pooling
        polling = torch.nn.AdaptiveAvgPool2d(1)
        tensor_pooled = polling(tensor)
        print("Polling shape ", tensor_pooled.shape)
        # sorting the averaged pooled values and getting the indexes 
        # print("Tensor-pooled ", tensor_pooled)
        sorted_tensor_pooled, sorted_indx = torch.sort(tensor_pooled, 1, descending=True)
        sorted_indx = sorted_indx.squeeze()

        a, b, c, d = tensor.shape
        mults = np.arange(0,a*b,b).reshape(-1,1)
        temp_indices = (mults+sorted_indx).flatten()
        temp = tensor.reshape(a*b,c,d)
        final = temp[temp_indices,:,:]
        sorted_tensor = final.reshape(a,b,c,d)
        return sorted_tensor

Hope it helps.!