Keras custom loss function huber

I wrote huber loss using Keras backend functions and this works well:

def huber_loss(y_true, y_pred, clip_delta=1.0):
    error = y_true - y_pred
    cond  = K.abs(error) < clip_delta
         
    squared_loss = 0.5 * K.square(error)
    linear_loss  = clip_delta * (K.abs(error) - 0.5 * clip_delta)

    return tf_where(cond, squared_loss, linear_loss)

But I need a more complex loss function:

1. if error <= A, use squared_loss
2. if A <= error < B, use linear_loss
3. if error >= B, used sqrt_loss

I wrote smth like this:

def best_loss(y_true, y_pred, A, B):
    error = K.abs(y_true - y_pred)
    cond  = error <= A
    cond2 = tf_logical_and(A < error, error <= B)

    squared_loss = 0.5 * K.square(error)
    linear_loss  = A * (error - 0.5 * A)
    sqrt_loss = A * np.sqrt(B) * K.sqrt(error) - 0.5 * A**2

    return tf_where(cond, squared_loss, tf_where(cond2, linear_loss, sqrt_loss))

But it does not work, the model with this loss function does not converge, what is the bug?