How to Fix CUDA Out of Memory in PyTorch

Q: How to fix CUDA out of memory in PyTorch?

Fix CUDA out of memory in order of effort: 1) Reduce batch size — the simplest fix, halving batch size roughly halves memory. 2) Use mixed precision with torch.cuda.amp.autocast() to save 30-40% memory. 3) Add gradient checkpointing with torch.utils.checkpoint to trade compute for memory. 4) For inference, wrap in torch.no_grad() to skip storing activations. 5) Use gradient accumulation to simulate large batches with small actual batches.

Reduce batch size first. Then try: torch.cuda.amp.autocast() for mixed precision, gradient checkpointing, or torch.no_grad() for inference.

Fix 1: Reduce Batch Size (Easiest)

# Halving batch size roughly halves activation memory
train_loader = DataLoader(dataset, batch_size=16)  # was 32

Fix 2: Mixed Precision Training (30-40% savings)

scaler = torch.cuda.amp.GradScaler()

for x, y in train_loader:
    optimizer.zero_grad()
    with torch.cuda.amp.autocast():  # uses float16 for activations
        output = model(x)
        loss = criterion(output, y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Fix 3: Gradient Checkpointing (50-80% activation savings)

from torch.utils.checkpoint import checkpoint

class MyModel(nn.Module):
    def forward(self, x):
        # Recompute activations during backward instead of storing
        x = checkpoint(self.block1, x, use_reentrant=False)
        x = checkpoint(self.block2, x, use_reentrant=False)
        return self.head(x)

Fix 4: torch.no_grad() for Inference

# Saves ~60% memory by not storing computation graph
with torch.no_grad():
    output = model(input_tensor)

Fix 5: Gradient Accumulation

# Simulate batch_size=32 with actual batch_size=8
accumulation_steps = 4
for i, (x, y) in enumerate(train_loader):
    output = model(x)  # batch_size=8
    loss = criterion(output, y) / accumulation_steps
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Quick Memory Check

# See how much VRAM you're using
print(f"Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
print(f"Cached: {torch.cuda.memory_reserved() / 1e9:.2f} GB")

# Force clear cache (rarely helps, but worth trying)
torch.cuda.empty_cache()