def wing_loss(
    output: torch.Tensor, target: torch.Tensor, width=5, curvature=0.5, reduction="mean"
):
    """Wing loss

    References:
        https://arxiv.org/pdf/1711.06753.pdf

    """
    diff_abs = (target - output).abs()
    loss = diff_abs.clone()

    idx_smaller = diff_abs < width
    idx_bigger = diff_abs >= width

    loss[idx_smaller] = width * torch.log(1 + diff_abs[idx_smaller] / curvature)

    C = width - width * math.log(1 + width / curvature)
    loss[idx_bigger] = loss[idx_bigger] - C

    if reduction == "sum":
        loss = loss.sum()

    if reduction == "mean":
        loss = loss.mean()

    return loss
Wing loss

    References:
        https://arxiv.org/pdf/1711.06753.pdf

    
segmentation_models_pytorch/losses/_functional.py