Direct Technical Solution: torch.cuda.empty_cache() releases only cached (unallocated) blocks back to the CUDA driver; it never frees memory occupied by active tensors (weights, optimizer states, computation graph nodes). Calling it inside your training loop hurts performance because CUDA must constantly re-allocate OS memory via costly system calls.

1. Root Cause: PyTorch Allocator Memory Fragmentation

Look closely at your error message: 18.21 GiB allocated + 4.80 GiB reserved. PyTorch had nearly 5 GB of memory held in its internal caching allocator, but it was split into scattered, non-contiguous memory chunks. When a tensor required 512 MiB of contiguous VRAM, the allocator failed to find a single chunk large enough.

2. The Modern Fix: Expandable Segments (PyTorch 2.0+)

The definitive solution in modern PyTorch is activating virtual memory management via the expandable_segments flag. This instructs CUDA to map physical memory pages to a contiguous virtual memory space, virtually eliminating memory fragmentation:

# Terminal / Docker Entrypoint (Set before launching Python)
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# Or directly in Python before calling any CUDA operations:
import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
import torch

3. The 4 Golden Rules to Prevent CUDA OOM

• Detach Loss Values: Never accumulate raw tensor losses: total_loss += loss retains the entire backward computation graph in VRAM! Always use total_loss += loss.item().
• Use Automatic Mixed Precision (AMP): Halve activation memory using native BF16/FP16:

  with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
      outputs = model(inputs)
      loss = criterion(outputs, targets)

• Gradient Accumulation: Instead of a batch size of 64 that OOMs, use a micro-batch size of 16 and accumulate gradients across 4 backward steps:

  loss = loss / 4
  loss.backward()
  if (step + 1) % 4 == 0:
      optimizer.step()
      optimizer.zero_grad(set_to_none=True)

• Zero Gradients with `set_to_none=True`: optimizer.zero_grad(set_to_none=True) deallocates memory instead of zeroing tensors with zeros of equal size.