The error looks like this:
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB
(GPU 0; 23.69 GiB total capacity; 20.15 GiB already allocated;
1.17 GiB free; 21.34 GiB reserved in total by PyTorch)
If reserved but unallocated memory is large try setting
max_split_size_mb to avoid fragmentation.Your training loop crashes, the process usually keeps the GPU reserved (leaving others unable to use it), and restarting the script often hits the same error immediately. Here is why it happens and five concrete fixes.
PyTorch allocates GPU memory for:
Memory compounds fast. A model with 100M parameters in float32 uses 400 MB for weights, 400 MB for gradients, and 800 MB for Adam state — before a single batch touches GPU. Large batches or deep networks with many intermediate activations push past GPU limits.
The fastest fix. Batch size scales linearly with activation memory. Halve the batch size, roughly halve the activation memory.
from torch.utils.data import DataLoader
# Before: OOM
train_loader = DataLoader(dataset, batch_size=64, shuffle=True)
# Fix: reduce batch size
train_loader = DataLoader(dataset, batch_size=16, shuffle=True)
If you need the effective batch size to stay large for training stability, use gradient accumulation — run multiple small batches before calling optimizer.step():
accumulation_steps = 4 # effective batch size = batch_size * accumulation_steps
optimizer.zero_grad()
for i, (inputs, labels) in enumerate(train_loader):
inputs, labels = inputs.cuda(), labels.cuda()
outputs = model(inputs)
loss = criterion(outputs, labels)
# Scale loss to account for accumulation
loss = loss / accumulation_steps
loss.backward()
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
torch.cuda.empty_cache() is widely misunderstood. It does not free memory that PyTorch is still using — it only releases memory from PyTorch's internal allocator cache back to CUDA. It will not fix an OOM that happens during a forward pass.
import torch
import gc
# After a crash or between training runs, do both:
gc.collect()
torch.cuda.empty_cache()
# Check what is actually using memory
print(torch.cuda.memory_summary())
The most common cause of unexpectedly high memory is holding references to tensors across iterations. This pattern silently accumulates memory every step:
# BAD: loss is a tensor on GPU, appending it keeps the computation graph alive
losses = []
for batch in train_loader:
loss = compute_loss(batch)
losses.append(loss) # holds graph in memory!
# GOOD: extract the scalar value
losses = []
for batch in train_loader:
loss = compute_loss(batch)
losses.append(loss.item()) # detaches from graph
During inference and validation, you do not need gradients. Without torch.no_grad(), PyTorch builds the full computation graph and stores all intermediate activations — doubling or tripling memory use for no reason.
import torch
model.eval()
# BAD: still builds computation graph
for inputs, labels in val_loader:
outputs = model(inputs.cuda())
loss = criterion(outputs, labels.cuda())
# GOOD: disables gradient tracking entirely
with torch.no_grad():
for inputs, labels in val_loader:
outputs = model(inputs.cuda())
loss = criterion(outputs, labels.cuda())
torch.no_grad() can also be used as a decorator on inference functions:
@torch.no_grad()
def predict(model, inputs):
model.eval()
return model(inputs.cuda())
Gradient checkpointing trades compute for memory. Instead of storing all intermediate activations during the forward pass (needed for backprop), it stores only a subset of checkpoints and recomputes the intermediate values during the backward pass. This typically reduces activation memory by ~sqrt(n_layers) at the cost of ~33% more compute.
import torch
import torch.utils.checkpoint as checkpoint
class CheckpointedModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.layers = torch.nn.ModuleList([
torch.nn.Linear(512, 512) for _ in range(12)
])
self.relu = torch.nn.ReLU()
def forward(self, x):
for layer in self.layers:
# Recompute activations on backward instead of storing them
x = checkpoint.checkpoint(lambda inp: self.relu(layer(inp)), x)
return x
For HuggingFace Transformers models, gradient checkpointing is a single flag:
from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")
model.gradient_checkpointing_enable() # one line
model.cuda()
Float32 tensors use 4 bytes per element. Float16 uses 2 bytes. Mixed precision training runs the forward pass in float16, which roughly halves the activation memory, then scales gradients back to float32 for the optimizer update to preserve numerical stability.
import torch
from torch.cuda.amp import autocast, GradScaler
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
scaler = GradScaler() # handles gradient scaling to avoid float16 underflow
for inputs, labels in train_loader:
inputs, labels = inputs.cuda(), labels.cuda()
optimizer.zero_grad()
# Forward pass in float16
with autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
# Backward pass with scaled gradients
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
Mixed precision is compatible with gradient checkpointing and gradient accumulation. Combining all three gives you the maximum memory reduction with acceptable compute overhead.
Before applying fixes blindly, profile to understand where your memory is going:
import torch
# Snapshot before and after a training step
torch.cuda.reset_peak_memory_stats()
# ... run your training step ...
print(torch.cuda.memory_summary(device=None, abbreviated=True))
# Key stats:
print(f"Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
print(f"Reserved: {torch.cuda.memory_reserved() / 1e9:.2f} GB")
print(f"Peak: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB")
To find which line of code triggers the OOM, use the memory snapshot tool introduced in PyTorch 2.0:
import torch
torch.cuda.memory._record_memory_history(max_entries=100000)
try:
# ... your training code ...
pass
except torch.cuda.OutOfMemoryError:
torch.cuda.memory._dump_snapshot("oom_snapshot.pickle")
raise
Load the snapshot in the PyTorch memory visualizer at pytorch.org/memory_viz to see a timeline of every allocation.
torch.no_grad() to your validation loop — zero downside, immediate gain..item() missing on accumulated tensors — common silent leak.autocast — ~2x memory reduction for activation tensors.For vector search and embedding workloads running alongside training, see FAISS index errors and memory patterns for GPU/CPU memory separation strategies.