import torch


def smooth_l1_loss_augmix(pred, target, beta=1.0):
    """Smooth L1 loss.

    Args:
        pred (torch.Tensor): The prediction.
        target (torch.Tensor): The learning target of the prediction.
        beta (float, optional): The threshold in the piecewise function.
            Defaults to 1.0.

    Returns:
        torch.Tensor: Calculated loss
    """
    pred_orig, _, _ = torch.chunk(pred, 3)
    target, _, _ = torch.chunk(target, 3)

    assert beta > 0
    if target.numel() == 0:
        return pred_orig.sum() * 0

    assert pred_orig.size() == target.size()
    diff = torch.abs(pred_orig - target)
    loss_orig = torch.where(diff < beta, 0.5 * diff * diff / beta,
                       diff - 0.5 * beta)

    return loss_orig