def lovasz_softmax_flat(probas, labels, classes='present'):
    """
    Multi-class Lovasz-Softmax loss.

    Args:
        probas (Tensor): Shape is [P, C], class probabilities at each prediction (between 0 and 1).
        labels (Tensor): Shape is [P], ground truth labels (between 0 and C - 1).
        classes (str|list): 'all' for all, 'present' for classes present in labels, or a list of classes to average.
    """
    if probas.numel() == 0:
        # only void pixels, the gradients should be 0
        return probas * 0.
    C = probas.shape[1]
    losses = []
    classes_to_sum = list(range(C)) if classes in ['all', 'present'
                                                   ] else classes
    for c in classes_to_sum:
        fg = paddle.cast(labels == c, probas.dtype)  # foreground for class c
        if classes == 'present' and fg.sum() == 0:
            continue
        fg.stop_gradient = True
        if C == 1:
            if len(classes_to_sum) > 1:
                raise ValueError('Sigmoid output possible only with 1 class')
            class_pred = probas[:, 0]
        else:
            class_pred = probas[:, c]
        errors = paddle.abs(fg - class_pred)
        if hasattr(paddle, "_legacy_C_ops"):
            errors_sorted, perm = paddle._legacy_C_ops.argsort(
                errors, 'axis', 0, 'descending', True)
        else:
            errors_sorted, perm = paddle._C_ops.argsort(errors, 'axis', 0,
                                                        'descending', True)
        errors_sorted.stop_gradient = False

        fg_sorted = paddle.gather(fg, perm)
        fg_sorted.stop_gradient = True

        grad = lovasz_grad(fg_sorted)
        grad.stop_gradient = True
        loss = paddle.sum(errors_sorted * grad)
        losses.append(loss)

    if len(classes_to_sum) == 1:
        return losses[0]

    losses_tensor = paddle.stack(losses)
    mean_loss = paddle.mean(losses_tensor)
    return mean_loss

    Multi-class Lovasz-Softmax loss.

    Args:
        probas (Tensor): Shape is [P, C], class probabilities at each prediction (between 0 and 1).
        labels (Tensor): Shape is [P], ground truth labels (between 0 and C - 1).
        classes (str|list): 'all' for all, 'present' for classes present in labels, or a list of classes to average.
    
paddleseg/models/losses/lovasz_loss.py