Skip to content

PyTorch GPU

In PyTorch, the GPU is mostly a device you move tensors and parameters onto. The model does not live on the GPU by magic; you put it there, move inputs there, and keep the whole computation on the same device.

This cheat sheet assumes CUDA is available. For availability and troubleshooting, see nvidia-cuda and cuda-troubleshooting.

Device Management

PyTorch devices look like strings:

import torch

cpu = torch.device("cpu")
gpu = torch.device("cuda")
gpu0 = torch.device("cuda:0")

Common device choices: - "cpu" — host CPU - "cuda" — default CUDA device - "cuda:0", "cuda:1" — specific GPU ordinal

Set the default device

torch.set_default_device("cuda")

Use this carefully. It affects tensor creation defaults, which can be convenient but can also hide device-mismatch bugs if you are not paying attention.

Check availability

print(torch.cuda.is_available())
print(torch.cuda.device_count())
print(torch.cuda.current_device())
print(torch.cuda.get_device_name(0))

If is_available() is False, stop and fix visibility/driver/toolkit before chasing model bugs.

Moving Tensors and Models to the GPU

The central operation is .to().

x = torch.randn(32, 128)
x_gpu = x.to("cuda")

model = MyModel().to("cuda")

A few rules: - Move inputs to the same device as the model. - Move labels and other operands to the right device too. - .cuda() is a shortcut; .to(device) is more flexible and preferred in most code. - .to() returns a new tensor; it does not always mutate in place.

Example:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

for x, y in loader:
    x, y = x.to(device), y.to(device)
    out = model(x)
    loss = loss_fn(out, y)
    ...

Mixed-device computations are a common source of errors. If one operand is on CPU and another on CUDA, many operations will fail or silently behave unexpectedly.

Tensor Creation on the GPU

Many factory functions accept a device argument:

x = torch.zeros(64, 128, device="cuda")
w = torch.randn(128, 10, device="cuda")
idx = torch.tensor([0, 1, 2], device=device)

If you create a tensor on CPU and then immediately move it, you have paid an extra copy for no reason. Prefer creating directly on the right device when practical.

model.to() vs model.cpu()

  • model.to("cuda") moves the model to GPU.
  • model.cpu() moves it back to CPU.
  • model.to(device) is the general form.

Use model.to(device) consistently so device logic is centralized.

Memory Basics

PyTorch on CUDA uses device memory for: - Model parameters - Buffers - Activations - Gradients - Optimizer states - Temporary tensors

Inspect memory

print(torch.cuda.memory_allocated())
print(torch.cuda.memory_reserved())
print(torch.cuda.max_memory_allocated())
  • allocated — currently used by tensors
  • reserved — cached allocator reservation, which can be larger than what is currently in use
  • max — peak so far

The caching allocator can make memory use look confusing. A tensor being out of scope does not always mean the memory is instantly returned to the system.

Clear cache

torch.cuda.empty_cache()

This releases unused cached memory back to the system. It does not free memory still in use, and it does not fix an out-of-memory problem caused by genuinely needing more memory.

Use empty_cache() when you want the system to see more free memory, not as a substitute for reducing memory pressure.

Autograd and the GPU

The GPU is where compute happens, but autograd still tracks operations.

Common points: - Backward pass happens on the same device as the forward computation for each tensor. - Gradients accumulate where parameters live. - Detach tensors you do not want tracked with .detach() or by using torch.no_grad().

with torch.no_grad():
    pred = model(x_gpu)

This is important for inference and evaluation, where you usually do not need gradients.

Precision and Autocast

Mixed precision is a common way to reduce memory and increase throughput on supported hardware.

scaler = torch.amp.GradScaler("cuda")

for x, y in loader:
    x, y = x.to(device), y.to(device)
    with torch.amp.autocast("cuda"):
        out = model(x)
        loss = loss_fn(out, y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Key ideas: - autocast selects lower-precision ops where safe. - A gradient scaler helps preserve numerical stability for training. - Mixed precision behavior is backend and precision dependent.

Use mixed precision when it helps, not blindly. For inference, autocast can also be useful, but the exact benefit depends on the model and hardware.

Common Device Errors

Error Typical cause
"Expected all tensors to be on the same device" operands on different devices
"CUDA out of memory" working set too large for the device
"All tensors must be on a CUDA device" some tensors are still on CPU
Silent slowness too many host-device transfers, tiny batches, no no_grad at inference
Device mismatch after saving/loading state dict loaded without matching device

Persisting and loading models are a frequent source of device surprises. A state dict does not strongly encode "this model is on CUDA." Device placement is separate from parameter values.

Safe load pattern

state = torch.load(path, map_location=device)
model.load_state_dict(state)
model.to(device)

Using map_location helps avoid loading everything onto CPU then unexpectedly expecting it to already be on GPU.

CPU vs CUDA Behavior

Not every operation behaves identically on CPU and CUDA. Watch for: - Different supported dtypes or backends - Different performance characteristics - Ops that fall back in surprising ways - Determinism differences

Code that works on CPU can still behave differently on GPU in edge cases, especially around precision and nondeterminism.

Multi-GPU Basics

If multiple GPUs are visible:

print(torch.cuda.device_count())
torch.cuda.set_device(0)

Common multi-GPU patterns: - Data parallelism — replicate the model, split the batch - Distributed training — processes, ranks, collective communication

torch.nn.DataParallel exists but is not the main path for serious multi-GPU training. For larger work, distributed approaches are more common. The details are framework- and task-specific.

A simple check before scaling up:

for i in range(torch.cuda.device_count()):
    print(i, torch.cuda.get_device_name(i))

Make sure your process can see the devices you think it can.

GPU Performance Checklist

  • Are inputs batched well for the device?
  • Are you moving data to CUDA outside the hot loop where possible?
  • Are you avoiding unnecessary .cpu() calls in training or inference loops?
  • Are you using torch.no_grad() for inference and evaluation?
  • Are you pinning memory and using async transfers where appropriate in dataloaders?
  • Are you watching memory use while the job runs?
  • Are you accidentally synchronizing often?

Small host-device transfer habits can dominate runtime more than the model itself.

Practical Inference Pattern

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model.eval()

with torch.no_grad():
    for x in batches:
        x = x.to(device)
        out = model(x)
        # process out

Keep the model on the device, move each batch there, and avoid tracking gradients during inference.

Practical Training Pattern

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(num_epochs):
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        out = model(x)
        loss = loss_fn(out, y)
        loss.backward()
        optimizer.step()

This is the minimal "everything lives on the GPU" training loop. Once this is solid, add mixed precision, gradient accumulation, checkpointing, or distributed training as needed.

Common Mistakes

  • Creating tensors on CPU and moving them repeatedly in a hot loop
  • Forgetting to move labels to the model's device
  • Using GPU tensors in host-only code paths by accident
  • Not using no_grad during inference, wasting memory and time
  • Assuming .cuda() everywhere is the same as managing one device variable
  • Ignoring memory reserved vs allocated when debugging OOM
  • Saving/loading without thinking about device placement
  • Hitting OOM and then calling empty_cache() instead of reducing memory use
  • gpu-basics — why GPU memory and parallelism matter
  • nvidia-cuda — CUDA stack context
  • cuda-troubleshooting — what to do when CUDA disappears or OOMs
  • llm-hardware — GPU memory in the context of large model inference and fine-tuning
  • mlops — reproducible device behavior across environments

Official Documentation

  • PyTorch CUDA notes and device semantics
  • torch.cuda API docs
  • torch.amp autocast and GradScaler docs
  • Your framework's GPU memory and multi-GPU notes

If a PyTorch GPU problem looks mysterious, the first question is usually device placement, and the second is usually memory.

Last Updated

2026-09-06 — Initial version.