Introduction

If adding @tf.function to a custom loss changes training results dramatically, the problem is usually not the decorator itself. The real issue is that graph execution exposes hidden assumptions in the loss, such as Python side effects, NumPy calls, frozen control flow, or state that was only behaving by accident in eager mode.

What @tf.function Changes

Without @tf.function, TensorFlow executes operations eagerly, one line at a time. That feels like normal Python and is easy to debug. With @tf.function, TensorFlow traces the function and builds a graph. Later calls reuse that graph instead of rerunning Python directly.

That shift matters because Python behavior and TensorFlow behavior are not identical:

• Python values can be captured at trace time
• Python side effects may run only during tracing
• NumPy operations execute outside the TensorFlow graph
• control flow may be converted or frozen depending on what is traceable

Keras already compiles parts of training internally, so decorating the loss can introduce another graph boundary. If the loss is not purely tensor-based, the extra compilation step can expose bugs immediately.

A Common Source Of Wrong Results

The loss function should depend on tensors, not Python state that can silently change between batches. This example shows a bad pattern:

1import tensorflow as tf
2
3use_absolute_error = True
4
5@tf.function
6def bad_loss(y_true, y_pred):
7    if use_absolute_error:
8        return tf.reduce_mean(tf.abs(y_true - y_pred))
9    return tf.reduce_mean(tf.square(y_true - y_pred))
10
11y_true = tf.constant([1.0, 2.0, 3.0])
12y_pred = tf.constant([1.5, 2.5, 2.0])
13
14print(bad_loss(y_true, y_pred).numpy())
15
16use_absolute_error = False
17print(bad_loss(y_true, y_pred).numpy())

This surprises many people. The Python if can be fixed when the function is traced, so changing use_absolute_error later may not produce the behavior you expect. The graph keeps using the traced branch until retracing happens.

In eager mode, the same code appears to work because Python reevaluates the if each call.

Keep The Loss Pure And Tensor-Based

A safer custom loss uses only TensorFlow ops and its explicit inputs.

1import tensorflow as tf
2
3def stable_loss(y_true, y_pred):
4    error = y_pred - y_true
5    return tf.reduce_mean(tf.square(error))
6
7model = tf.keras.Sequential([
8    tf.keras.layers.Input(shape=(1,)),
9    tf.keras.layers.Dense(1)
10])
11
12model.compile(
13    optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
14    loss=stable_loss
15)

This version is graph-friendly because it does not depend on external Python state, lists, counters, or NumPy arrays created inside the loss. Keras can wrap it as needed during model.fit.

In most cases, start here: write a correct tensor-only loss first, and let Keras decide when to trace it.

@tf.function Belongs More Naturally On The Training Step

If you need graph speedups, place @tf.function on the training step rather than on a fragile loss implementation.

1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(1,)),
5    tf.keras.layers.Dense(1)
6])
7
8optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)
9
10def mse_loss(y_true, y_pred):
11    return tf.reduce_mean(tf.square(y_true - y_pred))
12
13@tf.function
14def train_step(x, y):
15    with tf.GradientTape() as tape:
16        predictions = model(x, training=True)
17        loss_value = mse_loss(y, predictions)
18    gradients = tape.gradient(loss_value, model.trainable_variables)
19    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
20    return loss_value

This pattern is easier to reason about. The loss stays a simple mathematical function, while the expensive repeated step gets compiled.

Other Reasons Results Can Drift

Even if the loss looks clean, results can still change if it includes behavior that is sensitive to tracing:

• random numbers generated with Python or NumPy instead of TensorFlow ops
• mutable Python containers updated inside the loss
• calls to .numpy() inside graph-traced code
• shape-dependent branches that retrace unpredictably
• mixed precision or dtype conversions that differ between code paths

For example, a NumPy random sample inside a traced function may be computed during tracing instead of every step. That turns a dynamic loss term into a constant, which can absolutely change optimization behavior.

If you need randomness in the graph, use TensorFlow randomness:

noise = tf.random.normal(shape=tf.shape(y_pred), stddev=0.01)

If you need metrics, logging, or counters, keep them outside the loss or implement them with TensorFlow primitives designed for graph execution.

How To Debug The Difference

When a loss behaves differently under @tf.function, reduce the problem:

1. run the loss eagerly on fixed inputs
2. run the same loss under @tf.function on the same inputs
3. compare outputs before involving the optimizer or model training

If the outputs already differ, the bug is inside the loss. If they match, the difference is probably elsewhere in the training step, such as data shuffling, randomness, metric updates, or optimizer state.

During debugging, forcing eager execution can help:

tf.config.run_functions_eagerly(True)

Use that only as a temporary diagnostic tool. It is helpful for isolating graph-related assumptions, not as a permanent performance setting.

Common Pitfalls

• Decorating the loss first instead of making it tensor-pure first.
• Using Python if statements that depend on changing external values.
• Calling NumPy or .numpy() inside code expected to run as a graph.
• Mixing bookkeeping side effects with the mathematical loss computation.
• Assuming model.fit and a custom loop will behave identically if the loss contains hidden Python state.

Summary

• '@tf.function changes execution from eager Python to a traced TensorFlow graph.'
• Large result changes usually mean the loss relied on Python behavior that does not translate cleanly to graph execution.
• Write custom losses as pure tensor functions with no hidden external state.
• Put @tf.function on the training step when you want performance and clearer control.
• Compare eager and traced outputs on fixed inputs to find the exact source of divergence.