CPU, CUDA, MPS — Devices and Movement

One framework, three (or four) backends

PyTorch supports CPU, NVIDIA CUDA, and Apple Silicon MPS as first-class devices. ROCm (AMD) and XPU (Intel) exist but are less commonly seen in indie work. The mental model: a tensor lives on exactly one device; ops require their inputs to live on the same device; .to(device) moves data.

The device-string idiom

Almost every PyTorch project starts with a single line picking the best available device, then reuses that string everywhere:

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

That string then gets passed to .to(device) on the model and every batch. The "expected cuda:0 got cpu" error always traces to one tensor that didn't get moved.

MPS — the Apple Silicon backend

MPS works for the vast majority of standard ops. There are still a handful of operators that fall back to CPU (a warning is emitted by default unless you set PYTORCH_ENABLE_MPS_FALLBACK=1). torch.compile support on MPS is improving but still less mature than CUDA. For peak Apple Silicon performance, MLX is often a better fit because it was designed for the unified memory architecture from day one.

Code

Device detection and the canonical idiom·python

import torch

print(torch.cuda.is_available())          # True if NVIDIA GPU
print(torch.backends.mps.is_available())  # True on Apple Silicon
print(torch.cuda.device_count())          # number of CUDA GPUs

device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device}")

Move tensors and models·python

import torch
import torch.nn as nn

device = "mps"  # or whatever you picked

# Tensors
x = torch.randn(32, 100)            # CPU
x = x.to(device)                     # moves
print(x.device)                      # mps:0

# Create directly on device — avoids the CPU→device copy
y = torch.randn(32, 100, device=device)

# Models — to() moves all parameters AND buffers
model = nn.Linear(100, 10).to(device)

# Inputs and weights MUST be on the same device
out = model(y)   # works
# out = model(torch.randn(32, 100))  # RuntimeError: device mismatch

Pinned memory + non_blocking — the GPU training speedup·python

import torch
from torch.utils.data import DataLoader

# In your DataLoader, set pin_memory=True
# (only meaningful when you'll move data to a CUDA GPU)
loader = DataLoader(dataset, batch_size=64, num_workers=4, pin_memory=True)

device = "cuda"
for x, y in loader:
    # non_blocking=True lets the CPU keep going while the copy queues
    x = x.to(device, non_blocking=True)
    y = y.to(device, non_blocking=True)
    # ...training step...

CPU, CUDA, MPS — Devices and Movement

One framework, three (or four) backends

The device-string idiom

MPS — the Apple Silicon backend

Code

External links

Exercise

Progress

Comments 0