Introduction

Writing a custom TensorFlow loss with an if or else sounds simple until the condition depends on tensors. The main distinction is this: a normal Python if works for ordinary Python configuration values, but tensor-dependent branching inside the loss must use TensorFlow ops such as tf.where or tf.cond so the computation stays inside the graph.

When a Plain Python if Is Fine

If the branch depends on a regular Python flag chosen before training starts, a normal if is completely fine.

1import tensorflow as tf
2
3
4def make_loss(use_mae: bool):
5    if use_mae:
6        return tf.keras.losses.MeanAbsoluteError()
7    return tf.keras.losses.MeanSquaredError()
8
9loss_fn = make_loss(use_mae=True)

This is configuration-time branching. TensorFlow is not choosing per-example behavior here. Your Python code is just returning one loss function or another.

When the Branch Depends on Tensor Values

Problems start when the decision depends on y_true, y_pred, or some other tensor. A Python if cannot inspect a whole tensor the way you might inspect a scalar Python boolean.

Instead, use tensor operations. tf.where is often the simplest tool for element-wise branching.

1import tensorflow as tf
2
3
4def custom_loss(y_true, y_pred):
5    error = y_true - y_pred
6    absolute_error = tf.abs(error)
7    squared_error = tf.square(error)
8
9    per_example = tf.where(
10        absolute_error < 1.0,
11        0.5 * squared_error,
12        absolute_error - 0.5,
13    )
14
15    return tf.reduce_mean(per_example)

This is basically a Huber-style piecewise loss. Each element chooses a branch based on its own tensor value, which is exactly what tf.where is for.

Using tf.cond for Whole-Tensor Branches

If the condition should choose between two larger computations based on one tensor boolean, use tf.cond.

1import tensorflow as tf
2
3
4def custom_loss(y_true, y_pred):
5    mean_error = tf.reduce_mean(tf.abs(y_true - y_pred))
6
7    return tf.cond(
8        mean_error > 1.0,
9        lambda: tf.reduce_mean(tf.square(y_true - y_pred)),
10        lambda: tf.reduce_mean(tf.abs(y_true - y_pred)),
11    )

This is different from tf.where. tf.where is usually for element-wise selection, while tf.cond chooses between two branches of computation based on one condition.

Plugging a Custom Loss Into a Keras Model

Once the loss returns a scalar tensor, you can use it with model.compile.

1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(8, activation="relu"),
5    tf.keras.layers.Dense(1)
6])
7
8model.compile(optimizer="adam", loss=custom_loss)

That is all Keras needs. The important part is that the loss uses TensorFlow operations throughout.

Per-Example Weighting With Conditions

A practical use of conditional loss logic is weighting certain examples more heavily.

1import tensorflow as tf
2
3
4def weighted_loss(y_true, y_pred):
5    error = tf.abs(y_true - y_pred)
6    weights = tf.where(y_true > 10.0, 2.0, 1.0)
7    return tf.reduce_mean(error * weights)

This lets the model care more about errors in one region of the target space without leaving TensorFlow's execution model.

Why Python if Fails on Tensors

A tensor often represents many values at once. Python expects a single true or false value in an if statement. That mismatch is why code such as this is wrong:

# Do not do this when y_true is a tensor.
if y_true > 10:
    ...

TensorFlow needs the branch to remain symbolic or vectorized so it can execute it correctly across batches and inside traced functions.

Common Pitfalls

A common mistake is using a Python if with tensor values inside a custom loss. That works only for normal Python booleans, not for batch tensors.

Another issue is forgetting to reduce the loss to a scalar. Keras expects a final scalar loss value unless you are intentionally returning per-example losses for special handling.

Developers also sometimes mix NumPy operations into the loss. That breaks differentiability and pulls the computation out of TensorFlow.

Finally, do not overcomplicate the control flow if a standard built-in loss already matches your need. Many custom if and else losses are reimplementations of Huber, weighted cross-entropy, or another existing function.

Summary

• Use a normal Python if only for configuration-time choices.
• Use tf.where for element-wise tensor-dependent branches.
• Use tf.cond when one tensor boolean chooses between two larger computations.
• Keep the entire loss in TensorFlow ops so gradients work correctly.
• Reduce the result to the scalar loss value your model should optimize.