Introduction

Printing from inside a Keras loss function is usually a debugging task, not something you want to keep in normal training code. The main rule is that regular Python print is often the wrong tool inside TensorFlow execution, so the reliable options are tf.print, eager debugging, or moving the logging into a callback or custom training step.

Why Plain print Often Disappoints

Keras loss functions run inside TensorFlow execution, which may be graph-based or compiled. In that context, a normal Python print can run at trace time, run fewer times than you expect, or not reflect per-batch values in the way you intended.

That is why TensorFlow provides tf.print, which is itself a TensorFlow op and works during execution.

Use tf.print Inside the Loss Function

A simple custom loss can print tensors during training like this:

1import tensorflow as tf
2from tensorflow import keras
3import numpy as np
4
5
6def debug_mse(y_true, y_pred):
7    loss = tf.reduce_mean(tf.square(y_true - y_pred))
8    tf.print("batch loss:", loss)
9    tf.print("y_true:", y_true, summarize=4)
10    tf.print("y_pred:", y_pred, summarize=4)
11    return loss
12
13
14x = np.random.rand(32, 3).astype("float32")
15y = np.random.rand(32, 1).astype("float32")
16
17model = keras.Sequential([
18    keras.layers.Input(shape=(3,)),
19    keras.layers.Dense(8, activation="relu"),
20    keras.layers.Dense(1),
21])
22
23model.compile(optimizer="adam", loss=debug_mse)
24model.fit(x, y, epochs=2, batch_size=8, verbose=0)

This is the standard answer when you truly want to observe tensors inside the loss calculation.

Use run_eagerly=True for Heavier Debugging

If you need to step through Python logic more directly, compile the model with eager execution enabled.

1model.compile(
2    optimizer="adam",
3    loss=debug_mse,
4    run_eagerly=True,
5)

With eager execution, debugging becomes easier, but training is usually slower. This is a debugging mode, not the best production training configuration.

Consider a Callback Instead

Many times, printing inside the loss function is a sign that the real goal is to inspect per-batch or per-epoch values. A callback is often cleaner.

1class LossPrinter(keras.callbacks.Callback):
2    def on_train_batch_end(self, batch, logs=None):
3        print("batch", batch, "loss", logs.get("loss"))
4
5
6model.compile(optimizer="adam", loss="mse")
7model.fit(x, y, epochs=2, batch_size=8, callbacks=[LossPrinter()], verbose=0)

This avoids coupling debug output to the math of the loss function itself.

Custom Training Step for Advanced Inspection

If you want to inspect gradients, intermediate activations, or multiple loss components, a custom train_step is often the better design.

1class DebugModel(keras.Model):
2    def train_step(self, data):
3        x_batch, y_batch = data
4        with tf.GradientTape() as tape:
5            y_pred = self(x_batch, training=True)
6            loss = self.compiled_loss(y_batch, y_pred)
7        tf.print("train_step loss:", loss)
8        grads = tape.gradient(loss, self.trainable_variables)
9        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
10        self.compiled_metrics.update_state(y_batch, y_pred)
11        return {m.name: m.result() for m in self.metrics}

That approach is more flexible than forcing all debugging into the loss function.

Common Pitfalls

The most common mistake is using Python print and expecting it to behave like per-batch TensorFlow execution output.

Another issue is printing too much data and slowing training dramatically, especially with large tensors or many batches.

A third problem is debugging inside the loss function when the cleaner solution is actually a callback or a custom train_step.

Summary

• Use tf.print instead of Python print inside TensorFlow loss code.
• Turn on run_eagerly=True when you need easier debugging and can tolerate slower execution.
• Prefer callbacks when you only need batch or epoch loss values.
• Use a custom train_step for deeper inspection of training behavior.
• Treat printing inside the loss function as temporary debugging, not as normal training design. If you need the same visibility for every experiment, it usually belongs in callbacks, metrics, or a custom training loop instead. That makes the debug signal easier to turn on and off without touching the loss math itself.