Introduction

A TensorFlow job that runs for a few epochs and then suddenly fails with GPU out-of-memory is usually not just "the model is too large." If the same model fits at the beginning but crashes later, memory is often accumulating somewhere across steps or epochs. The real fix is to find what is growing over time: tensors kept alive in Python, retraced graphs, cached dataset state, or large activations combined with an aggressive batch size.

Why OOM Can Appear Late Instead Of Immediately

If the model is simply too big for the GPU, you usually fail on the first forward or backward pass. A delayed OOM suggests memory growth.

Common causes include:

• storing predictions or losses in a Python list every step
• recreating @tf.function graphs repeatedly with changing input shapes
• building new models inside a loop without clearing the old one
• dataset pipelines that cache large tensors unexpectedly
• fragmentation or aggressive preallocation on the GPU

That is why "it crashes after epoch 6" is an important clue. Something is not being reused or released correctly.

Enable Safer GPU Memory Behavior

A good first step is to tell TensorFlow not to grab all GPU memory up front.

1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices('GPU')
4for gpu in gpus:
5    tf.config.experimental.set_memory_growth(gpu, True)

This does not solve every OOM problem, but it makes memory use more incremental and can reduce confusion while you debug.

Avoid Accumulating Tensors Across Epochs

One of the easiest ways to leak memory is to store tensors every step without converting or discarding them.

Problematic pattern:

1history_tensors = []
2
3for batch_x, batch_y in dataset:
4    with tf.GradientTape() as tape:
5        preds = model(batch_x, training=True)
6        loss = loss_fn(batch_y, preds)
7    history_tensors.append(loss)

loss here is still a TensorFlow tensor, and holding onto many of them can keep computation history or device memory alive longer than expected.

Safer version:

1history_values = []
2
3for batch_x, batch_y in dataset:
4    with tf.GradientTape() as tape:
5        preds = model(batch_x, training=True)
6        loss = loss_fn(batch_y, preds)
7    history_values.append(float(loss.numpy()))

If you only need scalar logging, store plain Python numbers instead of tensors.

Watch Out For Repeated Model Creation

Another common mistake is creating a new model in a loop, especially during experiments or cross-validation, without clearing the previous graph state.

1import tensorflow as tf
2
3for run in range(3):
4    tf.keras.backend.clear_session()
5    model = tf.keras.Sequential([
6        tf.keras.layers.Dense(128, activation='relu'),
7        tf.keras.layers.Dense(10)
8    ])
9    model.compile(optimizer='adam', loss='mse')
10    model.fit(x_train, y_train, epochs=2)

clear_session() is important when many model objects are created in the same Python process.

Reduce Peak Memory Pressure

Even when the late OOM is caused by growth, peak usage still matters. The most direct levers are:

• lower the batch size
• reduce image or sequence dimensions
• use mixed precision when appropriate
• simplify the model or checkpoint fewer activations

Mixed precision can help on supported GPUs:

from tensorflow.keras import mixed_precision

mixed_precision.set_global_policy('mixed_float16')

That often lowers activation memory significantly, though you should still validate numerical behavior.

Dataset Pipelines Matter Too

A tf.data pipeline can also cause problems if it caches very large tensors in memory or expands data aggressively.

dataset = dataset.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)

This is usually sensible. But cache() should be used carefully. Caching a large transformed dataset in memory can make a job appear fine at first and then run out once the cache fills.

Common Pitfalls

The most common mistake is treating a late OOM exactly like an immediate OOM. If the crash happens only after several epochs, look for growth, not just peak size.

Another issue is storing tensors for logging or debugging instead of converting them to plain values. That keeps extra memory alive longer than needed.

It is also easy to rebuild models repeatedly in notebooks and long-running processes without clearing the old state.

Finally, do not forget the input pipeline. A dataset cache or oversized prefetch configuration can contribute to memory pressure just as much as the model itself.

Summary

• A delayed TensorFlow GPU OOM usually points to memory growth over time, not only a large model.
• Check for stored tensors, repeated graph creation, and problematic dataset caching.
• Enable GPU memory growth to make debugging behavior clearer.
• Use clear_session() when creating many models in one process.
• Lower batch size, use mixed precision, and simplify the input pipeline when peak memory is still too high.