Solved by Kushaj in post #2 
              
              
                def func(x):
    _, indices = torch.topk(x, k=x.size(1), dim=1)
    _, indices = torch.sort(x, dim=1)
    mask = x <= 1292
    num_greater = mask.sum(dim=1, keepdim=True)
    
    _, s_indices = torch.sort(indices, dim=1)
    s_indices -= num_greater - 1
    
    return torch.where(mask, x, s_indice…