PyTorch GPU: CUDA and MPS

Running PyTorch on a GPU can make training 10-50 times faster than a CPU, because neural network math is massively parallel. PyTorch supports NVIDIA GPUs through CUDA, AMD GPUs through ROCm (which also uses the "cuda" device name), and Apple Silicon through MPS.

Picking a Device

import torch

if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")
print(device)

cuda   (sample — mps on a Mac, cpu otherwise)

Recent PyTorch versions also offer torch.accelerator.current_accelerator(), which returns whichever accelerator is present, but the explicit version above works everywhere and is easy to read.

Moving Models and Data

Both the model and every batch must be on the same device. .to(device) moves a model in place, but returns a new tensor for tensors, so assign the result.

from torch import nn

model = nn.Linear(10, 2).to(device)
x = torch.randn(4, 10).to(device)
print(model(x).device)

# create directly on the device
z = torch.zeros(3, device=device)

cuda:0   (sample)

The Device Mismatch Error

The most common GPU error looks like this:

x_cpu = torch.randn(4, 10)
model(x_cpu)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Fix it by moving inputs with .to(device) inside your training loop. To bring results back for NumPy or plotting, use tensor.cpu().

Checking GPU Memory

if device.type == "cuda":
    print(torch.cuda.get_device_name(0))
    print(round(torch.cuda.memory_allocated() / 1e6, 1), "MB allocated")

NVIDIA GeForce RTX 4070
0.1 MB allocated
(sample output)

If you see "CUDA out of memory", reduce the batch size, use mixed precision, or free tensors you no longer need. nvidia-smi in a terminal shows overall GPU usage.

Mixed Precision Training

Automatic mixed precision runs most operations in 16-bit, which is faster and uses less memory, while keeping sensitive operations in 32-bit. bfloat16 is the simplest on modern GPUs:

opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
xb, yb = torch.randn(32, 10, device=device), torch.randint(0, 2, (32,), device=device)

with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
    loss = nn.functional.cross_entropy(model(xb), yb)
opt.zero_grad()
loss.backward()
opt.step()
print(round(loss.item(), 4))

0.7231   (sample output)

With float16 instead of bfloat16, also use torch.amp.GradScaler to prevent tiny gradients from underflowing to zero.

Speed Tips

• Use DataLoader(pin_memory=True) and .to(device, non_blocking=True) on CUDA.
• Avoid calling .item() or .cpu() every step; they force the GPU to synchronize.
• Try model = torch.compile(model) for extra speed on supported setups.
• Not every operation is implemented on MPS yet; setting PYTORCH_ENABLE_MPS_FALLBACK=1 runs missing ones on the CPU.

Frequently asked questions

How do I use the GPU in PyTorch?

Create device = torch.device("cuda") (or "mps" on a Mac), move the model with model.to(device), and move every batch with xb.to(device) inside the loop.

Does PyTorch work on Apple M-series Macs?

Yes. The standard macOS build includes the MPS backend, which uses the Apple GPU. Most common layers are supported.

How do I fix CUDA out of memory in PyTorch?

Lower the batch size, enable mixed precision, use gradient accumulation, and make sure you are not keeping references to old tensors such as un-detached losses.