Hot to fix Tensorflow model not running in Eager mode with .fit?

TensorFlow

Eager Execution

Model Training

Debugging

Machine Learning

Hot to fix Tensorflow model not running in Eager mode with .fit?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

TensorFlow & Keras Machine Learning Debugging & Errors

Introduction

By default, TensorFlow 2.x runs model.fit() in graph mode (via tf.function) for performance, even though eager execution is enabled globally. This means print() statements, pdb breakpoints, and NumPy operations inside your model's call() method do not execute as expected during training. To force eager execution during fit(), pass run_eagerly=True to model.compile(). This is essential for debugging custom layers, custom training steps, and complex model logic, but should be disabled for production training since graph mode is significantly faster.

The Problem

python

1import tensorflow as tf
2
3class MyModel(tf.keras.Model):
4    def __init__(self):
5        super().__init__()
6        self.dense = tf.keras.layers.Dense(10)
7
8    def call(self, inputs):
9        print(f"Input shape: {inputs.shape}")  # Does NOT print during fit()
10        x = self.dense(inputs)
11        print(f"Output: {x}")  # Does NOT print during fit()
12        return x
13
14model = MyModel()
15model.compile(optimizer='adam', loss='mse')
16
17# print() statements inside call() are silenced during fit()
18model.fit(x_train, y_train, epochs=5)
19# No print output — model.fit() runs in graph mode by default

Fix 1: run_eagerly=True in compile()

python

1model = MyModel()
2model.compile(optimizer='adam', loss='mse', run_eagerly=True)
3
4model.fit(x_train, y_train, epochs=5)
5# Now prints:
6# Input shape: (32, 784)
7# Output: tf.Tensor([...], shape=(32, 10), dtype=float32)

This forces fit() to run every operation eagerly, including call(), custom loss functions, and custom metrics.

Fix 2: tf.config.run_functions_eagerly(True)

Force all tf.function-decorated code to run eagerly:

python

1# Enable globally — affects all tf.function calls
2tf.config.run_functions_eagerly(True)
3
4model.compile(optimizer='adam', loss='mse')  # No run_eagerly needed
5model.fit(x_train, y_train, epochs=5)
6# print() works, pdb works, numpy operations work
7
8# Disable after debugging
9tf.config.run_functions_eagerly(False)

Fix 3: tf.debugging for Graph Mode

If you want graph-mode speed but still need debugging output:

python

1class MyModel(tf.keras.Model):
2    def call(self, inputs):
3        # tf.print works in both eager and graph mode
4        tf.print("Input shape:", tf.shape(inputs))
5        tf.print("Input mean:", tf.reduce_mean(inputs))
6
7        x = self.dense(inputs)
8
9        # Assert in graph mode
10        tf.debugging.assert_all_finite(x, "NaN detected in output")
11
12        return x
13
14model = MyModel()
15model.compile(optimizer='adam', loss='mse')  # Default graph mode
16model.fit(x_train, y_train, epochs=5)
17# tf.print output appears during training

Debugging Custom Layers

python

1class DebugLayer(tf.keras.layers.Layer):
2    def __init__(self, units):
3        super().__init__()
4        self.dense = tf.keras.layers.Dense(units)
5
6    def call(self, inputs, training=None):
7        x = self.dense(inputs)
8
9        # These only work in eager mode (run_eagerly=True)
10        if training:
11            print(f"Training mode - output stats: mean={x.numpy().mean():.4f}")
12
13        # These work in both modes
14        tf.print("Layer output shape:", tf.shape(x))
15
16        return x
17
18# For debugging
19model.compile(optimizer='adam', loss='mse', run_eagerly=True)
20model.fit(x_train, y_train, epochs=1)
21
22# For production (faster)
23model.compile(optimizer='adam', loss='mse', run_eagerly=False)
24model.fit(x_train, y_train, epochs=50)

Custom Training Step Debugging

python

1class CustomModel(tf.keras.Model):
2    def train_step(self, data):
3        x, y = data
4
5        with tf.GradientTape() as tape:
6            y_pred = self(x, training=True)
7            loss = self.compute_loss(y=y, y_pred=y_pred)
8
9        # Debug: inspect gradients
10        gradients = tape.gradient(loss, self.trainable_variables)
11
12        # Only works with run_eagerly=True
13        for g, v in zip(gradients, self.trainable_variables):
14            if g is not None:
15                print(f"{v.name}: grad_norm={tf.norm(g).numpy():.4f}")
16
17        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
18        return {m.name: m.result() for m in self.metrics}
19
20model = CustomModel()
21model.compile(optimizer='adam', loss='mse', run_eagerly=True)
22model.fit(x_train, y_train, epochs=1)

Checking Current Execution Mode

python

1# Check if eager mode is enabled globally
2print(tf.executing_eagerly())  # True (default in TF2)
3
4# Check inside a tf.function
5@tf.function
6def check_mode():
7    tf.print("Eager inside tf.function:", tf.executing_eagerly())
8    # Prints: Eager inside tf.function: False (unless run_functions_eagerly is True)
9
10check_mode()
11
12# Check run_eagerly setting on a compiled model
13model.compile(optimizer='adam', loss='mse', run_eagerly=True)
14print(model.run_eagerly)  # True

Performance Impact

python

1import time
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(256, activation='relu'),
5    tf.keras.layers.Dense(128, activation='relu'),
6    tf.keras.layers.Dense(10, activation='softmax')
7])
8
9# Graph mode (fast — default)
10model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', run_eagerly=False)
11start = time.time()
12model.fit(x_train, y_train, epochs=10, verbose=0)
13print(f"Graph mode: {time.time() - start:.2f}s")
14
15# Eager mode (slow — for debugging)
16model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', run_eagerly=True)
17start = time.time()
18model.fit(x_train, y_train, epochs=10, verbose=0)
19print(f"Eager mode: {time.time() - start:.2f}s")
20
21# Typical results:
22# Graph mode: 8.5s
23# Eager mode: 25.3s (3x slower)

Common Pitfalls

Leaving run_eagerly=True in production: Eager mode is 2-5x slower than graph mode because it cannot optimize the computation graph. Use run_eagerly=True only for debugging, then switch back to False (or omit it) for actual training. Consider using an environment variable to toggle: run_eagerly=os.environ.get('DEBUG', False).
Using print() instead of tf.print() in model code: Python print() only executes during graph tracing (once), not during actual training. Use tf.print() for output that should appear every step. tf.print is a TensorFlow operation that executes within the graph.
Calling .numpy() inside model code without eager mode: tensor.numpy() only works in eager mode. In graph mode, it raises AttributeError: 'Tensor' object has no attribute 'numpy'. Use tf.print or tf.debugging assertions instead, or enable run_eagerly=True.
Confusing tf.executing_eagerly() with model.run_eagerly: tf.executing_eagerly() returns True at the top level in TF2, but inside model.fit() the call() method runs inside a tf.function (graph mode) regardless. Only model.compile(run_eagerly=True) forces the training loop itself to be eager.
Custom loss functions not working as expected in graph mode: If your custom loss uses Python control flow (if, for with Python variables), it may behave differently in graph mode vs eager mode. Use tf.cond, tf.while_loop, or run_eagerly=True to ensure correct behavior, then convert to TF ops for production.

Summary

model.fit() runs in graph mode by default for performance, even though TF2 enables eager execution globally
Pass run_eagerly=True to model.compile() to force eager execution during training (for debugging)
Use tf.print() instead of print() and tf.debugging assertions for output that works in both modes
tf.config.run_functions_eagerly(True) forces all tf.function calls to run eagerly (global setting)
Always disable eager mode for production training — graph mode is 2-5x faster

Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course

Basis Lab

No course covers this one yet

Basis Lab writes one for you from a sentence about what you want to be able to do, then teaches it and checks you understood. Your first course is free.

Build my course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Hot to fix Tensorflow model not running in Eager mode with .fit?

ML System Design practice on Codemia

Introduction

The Problem

Fix 1: run_eagerly=True in compile()

Fix 2: tf.config.run_functions_eagerly(True)

Fix 3: tf.debugging for Graph Mode

Debugging Custom Layers

Custom Training Step Debugging

Checking Current Execution Mode

Performance Impact

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

No course covers this one yet

Track what you have practised

ML System Design practice on Codemia