# Debug Out-of-Memory Errors When Training Large Models in Jupyter Notebooks

> Debug out-of-memory errors training large models in Jupyter notebooks. Learn to fix GPU OOM crashes with TensorFlow memory growth, smaller batch sizes, and PyTorch CUDA cache clearing.

- Repository: [Microsoft/AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners)
- Tags: how-to-guide
- Published: 2026-08-01

---

**Enable TensorFlow memory growth, reduce batch sizes, and clear PyTorch CUDA caches to eliminate GPU OOM crashes while working through the Microsoft AI For Beginners curriculum.**

Training deep-learning models in the *AI-For-Beginners* notebooks can exceed the memory limits of your GPU or CPU, leading to out-of-memory (OOM) failures that halt experimentation. The curriculum demonstrates battle-tested patterns in `lessons/5-NLP/16-RNN/RNNTF.ipynb` and `lessons/4-ComputerVision/10-GANs/GANPyTorch.ipynb` that you can apply to any notebook to debug and resolve these errors.

## Why OOM Errors Occur in AI For Beginners Notebooks

The notebooks in the microsoft/AI-For-Beginners repository load full datasets into memory and instantiate large neural networks by default. When running on consumer GPUs with limited VRAM (8–12 GB), pre-allocated TensorFlow buffers or aggressive PyTorch DataLoader prefetching quickly exhaust available memory. The repository source code shows these issues most frequently in recurrent neural network (RNN) and generative adversarial network (GAN) lessons where batch processing demands are highest.

## Immediate Fixes Using Repository Patterns

The following strategies are extracted directly from the notebook implementations and configuration files in the repository.

### Enable TensorFlow Memory Growth

By default, TensorFlow allocates the entire GPU memory at startup, preventing other processes or larger batches from fitting. The RNN notebooks demonstrate how to enable dynamic memory growth.

Add this snippet at the top of any TensorFlow notebook:

```python
import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        tf.config.experimental.set_memory_growth(gpus[0], True)
    except RuntimeError as e:
        print(e)

```

This pattern appears in `lessons/5-NLP/16-RNN/RNNTF.ipynb` at line 56, where it prevents OOM errors during LSTM training.

### Reduce Batch Size

Smaller batches require less GPU RAM per optimization step. The curriculum varies batch sizes by model complexity.

- **NLP notebooks**: Use `batch_size = 16` as shown in `RNNTF.ipynb` at line 93.
- **GAN notebooks**: Use `batch_size = 256` in `GANPyTorch.ipynb` at line 70, though you should reduce this to 64 or 32 if OOM occurs.

Adjust the `batch_size` variable or the `batch_size` argument of the DataLoader:

```python

# PyTorch example from the repository patterns

from torch.utils.data import DataLoader

batch_size = 16  # Start small, increase if memory permits

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

```

### Clear PyTorch CUDA Cache

PyTorch retains unused memory in a cache for speed, which can trigger false OOM errors. Explicitly clear this cache after each epoch or when switching models.

Insert this after your validation loop:

```python
import torch

for epoch in range(num_epochs):
    # Training loop...

    torch.cuda.empty_cache()  # Free unused memory

    print(f'Epoch {epoch} finished')

```

### Limit Data Prefetching

The `prefetch` operation holds several batches in RAM simultaneously. Set a low prefetch factor or remove it entirely to reduce memory pressure.

Look for patterns like this in the TensorFlow notebooks:

```python
ds_train.map(...).batch(batch_size).prefetch(1)  # Change prefetch(1) to prefetch(0) or remove

```

### Enable Mixed Precision Training

Mixed precision halves the memory footprint of activations and weights by using float16 instead of float32 for certain operations.

Add this before model creation in TensorFlow notebooks:

```python
import tensorflow as tf

tf.keras.mixed_precision.set_global_policy('mixed_float16')
model = tf.keras.Sequential([...])  # Build model after this call

```

## Step-by-Step Debugging Workflow

Follow this systematic approach to diagnose and fix OOM errors in any AI For Beginners notebook:

1. **Check hardware availability** – Verify GPU detection by running `!nvidia-smi` in a notebook cell.
2. **Enable memory growth** – Insert the TensorFlow memory-growth snippet early in the notebook, before any model instantiation.
3. **Decrease batch size** – Halve the `batch_size` variable until the notebook runs without error.
4. **Add cache clearing** – Call `torch.cuda.empty_cache()` after each epoch in PyTorch notebooks.
5. **Monitor consumption** – Insert memory stats printing after training steps using `nvidia-smi` via subprocess.
6. **Iterate** – If OOM persists, enable mixed-precision or switch to `IterableDataset` for lazy loading.

## Key Files in the Repository

| File | Relevance |
|------|-----------|
| `lessons/5-NLP/16-RNN/RNNTF.ipynb` | Demonstrates TensorFlow memory-growth setup (`tf.config.experimental.set_memory_growth`) and batch-size handling at line 56. |
| `lessons/5-NLP/16-RNN/RNNPyTorch.ipynb` | Shows PyTorch DataLoader batch-size tuning and tensor device management. |
| `lessons/4-ComputerVision/10-GANs/GANPyTorch.ipynb` | Contains GAN training loops where batch size is the primary OOM lever (line 70). |
| [`lessons/0-course-setup/setup.md`](https://github.com/microsoft/AI-For-Beginners/blob/main/lessons/0-course-setup/setup.md) | Guides environment creation and CUDA version compatibility checks. |

## Summary

- **Enable memory growth** in TensorFlow using `tf.config.experimental.set_memory_growth` to prevent full GPU allocation at startup.
- **Reduce batch sizes** to 16 or lower for NLP tasks, and monitor `batch_size` variables in DataLoader configurations.
- **Clear PyTorch caches** with `torch.cuda.empty_cache()` after each epoch to release reserved but unused VRAM.
- **Use mixed precision** via `tf.keras.mixed_precision.set_global_policy('mixed_float16')` to halve memory usage.
- **Reference specific lines** in `RNNTF.ipynb` (line 56, 93) and `GANPyTorch.ipynb` (line 70) for working implementations.

## Frequently Asked Questions

### Why does TensorFlow crash with OOM even though nvidia-smi shows free memory?

TensorFlow pre-allocates the entire GPU memory by default, reserving it exclusively for the process. Enable memory growth with `tf.config.experimental.set_memory_growth(physical_devices[0], True)` as implemented in the RNN notebooks to allow dynamic allocation.

### How do I know if I should reduce batch size or switch to mixed precision?

Start with batch size reduction because it is framework-agnostic and immediately effective. If reducing the batch below 8 still causes OOM, implement mixed precision training to cut memory usage by approximately 50% without changing the batch dimension.

### Where should I place the torch.cuda.empty_cache() call in my training loop?

Place it after the validation phase or at the end of each epoch, outside the inner batch loop. This ensures you clear gradients and activations accumulated during the epoch without slowing down the training step itself.

### Can I run these notebooks without a GPU at all?

Yes, though training will be slower. TensorFlow automatically falls back to CPU, but for PyTorch notebooks you may need to explicitly set `device="cpu"` when loading models. Check [`lessons/0-course-setup/setup.md`](https://github.com/microsoft/AI-For-Beginners/blob/main/lessons/0-course-setup/setup.md) for CPU-compatible environment configurations.