TensorFlow
error handling
Python
exception handling
programming issues

Tensorflow, try and except doesn't handle exception

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

Introduction

When try and except appears to miss TensorFlow errors, the issue is usually execution timing. In TensorFlow graph-style code, many failures appear only when operations execute, not when they are defined. Correct handling depends on wrapping the real execution point and catching exception types that match runtime behavior.

Know Where TensorFlow Actually Raises

In deferred execution patterns, graph construction succeeds and errors appear later at run time.

python
1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5x = tf.compat.v1.placeholder(tf.float32, shape=[None, 2])
6y = tf.reshape(x, [3, 3])
7
8with tf.compat.v1.Session() as sess:
9    try:
10        sess.run(y, feed_dict={x: [[1.0, 2.0], [3.0, 4.0]]})
11    except Exception as e:
12        print("Caught:", type(e).__name__)

Wrapping graph definition alone would not catch this runtime shape failure.

Eager Execution Behaves Differently

In eager mode, many errors are immediate.

python
1import tensorflow as tf
2
3try:
4    x = tf.constant([[1, 2], [3, 4]])
5    y = tf.reshape(x, [3, 3])
6except Exception as e:
7    print("Caught eager error:", type(e).__name__)

This difference causes confusion when teams mix TensorFlow one and two patterns.

Catch Specific Exception Classes

Catching broad Exception works for debugging, but production code should catch specific classes where possible.

python
1from tensorflow.python.framework.errors_impl import InvalidArgumentError
2
3try:
4    # run operation
5    pass
6except InvalidArgumentError as e:
7    print("Invalid argument:", e)

Specific catches improve logging quality and recovery decisions.

Handle Dataset Errors at Iteration Time

Errors in tf.data pipelines often surface when iterating, not when defining map operations.

python
1import tensorflow as tf
2
3ds = tf.data.Dataset.from_tensor_slices([1, 2, 0])
4ds = ds.map(lambda x: 10 // x)
5
6try:
7    for item in ds:
8        print(item.numpy())
9except Exception as e:
10    print("Caught during iteration:", type(e).__name__)

Wrap the consumer loop when diagnosing dataset failures.

Understand tf.function Tracing and Runtime Errors

With @tf.function, exceptions can occur during tracing or execution depending on inputs.

python
1@tf.function
2def risky(x):
3    return tf.reshape(x, [3, 3])
4
5try:
6    out = risky(tf.constant([[1, 2], [3, 4]]))
7except Exception as e:
8    print("Caught:", type(e).__name__)

Call-site wrapping is still required.

Improve Diagnostics with Assertions

Use TensorFlow assertions to fail early with clear messages.

python
def safe_reshape(x):
    tf.debugging.assert_equal(tf.size(x) % 9, 0, message="Tensor size must be divisible by 9")
    return tf.reshape(x, [-1, 9])

Clear assertion messages reduce debugging time compared with opaque internal errors.

Handle Async and Callback Contexts

In training callbacks or async job runners, exceptions may not bubble to outer scope where try is located. Capture exceptions at callback boundaries and propagate them through explicit status handling.

This is important in serving and pipeline workers where silent callback failures can mask operational incidents.

Practical Debug Checklist

When exceptions are not caught as expected:

  1. Confirm eager versus graph execution mode.
  2. Move try around operation execution line.
  3. Inspect tf.data consumer loops.
  4. Add assertion guards before risky ops.
  5. Narrow exception class and improve logs.

Following this order usually isolates root cause quickly.

Rethrow with Context in Service Layers

In production services, it can be useful to catch low-level TensorFlow exceptions, attach request context, and rethrow domain-specific errors.

python
1try:
2    prediction = model_fn(batch)
3except Exception as exc:
4    raise RuntimeError(f\"Inference failed for batch_id={batch_id}\") from exc

This preserves root traceback while adding useful operational context for logs and alerts.

Common Pitfalls

  • Wrapping graph construction but not execution calls.
  • Catching very broad exceptions and hiding useful context.
  • Expecting dataset map definition to raise immediately.
  • Mixing TensorFlow one and TensorFlow two mental models.
  • Ignoring callback and asynchronous execution boundaries.

Summary

  • TensorFlow exception handling depends on execution timing.
  • Wrap actual operation execution, not only graph definition.
  • Treat eager, graph, and tf.function contexts differently.
  • Use specific exception classes and assertion guards.
  • Add boundary-level handling for dataset and async paths.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

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

Start the free 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

All Rights Reserved.