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_indices)

x = torch.tensor([[1296, 1295, 1292, 4, 1311, 4, 1293, 2],
                  [1297, 1295, 1292, 1401, 1294, 4, 1293, 2]], dtype=torch.int64)
func(x)

# tensor([[   3,    2, 1292,    4,    4,    4,    1,    2],
#         [   4,    3, 1292,    5,    2,    4,    1,    2]])