Solved by ptrblck in post #2 
              
              
                You could directly index the tensor as seen here: 
feature = torch.rand(4,6,3,3)
weight = torch.rand(4,6)
_, indices = torch.sort(weight, dim=1, descending = True)
z1 = feature[0,indices[0],:,:]
z2 = feature[1,indices[1],:,:]
z3 = feature[2,indices[2],:,:]
z4 = feature[3,indices[3],:,:]
z = torch.st…