How do you use GPUs in PyTorch and what are the key patterns for writing device-agnostic code?

PyTorch's device abstraction allows the same code to run on CPU, single GPU, or multiple GPUs with minimal changes. The fundamental operations are moving tensors to a device with .to(device) or .cuda(), and ensuring model and data tensors always reside on the same device before any computation.

A critical performance concept: CPU–GPU data transfers are expensive (PCIe bandwidth is limited vs. GPU memory bandwidth). Minimise them by loading data onto the GPU once per batch, pre-computing dataset statistics on CPU, and avoiding frequent tensor transfers inside the training loop.

import torch
import torch.nn as nn

# Device-agnostic code pattern
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using: {device}')  # cuda / mps / cpu

# Move model to device
model = nn.Linear(10, 5).to(device)

# Move data to device in the training loop
for X_batch, y_batch in loader:
    X_batch = X_batch.to(device, non_blocking=True)
    y_batch = y_batch.to(device, non_blocking=True)
    y_hat = model(X_batch)
    # ...

# Check which device a tensor is on
t = torch.randn(3)
print(t.device)         # cpu
t_gpu = t.cuda()        # or t.to('cuda:0')
print(t_gpu.device)     # cuda:0

# Apple Silicon
device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')

# Memory diagnostics
print(torch.cuda.memory_allocated() / 1e9, 'GB allocated')
print(torch.cuda.max_memory_allocated() / 1e9, 'GB peak')
torch.cuda.empty_cache()  # release unused cached GPU memory

# Multi-GPU: DistributedDataParallel (DDP) preferred over DataParallel
model_ddp = nn.parallel.DistributedDataParallel(model, device_ids=[0, 1])