CPU / GPU / MPS

import torch

# ===== Detect available device =====
device = (
    'cuda' if torch.cuda.is_available()
    else 'mps' if torch.backends.mps.is_available()    # Apple Silicon
    else 'cpu'
)
print('device:', device)

# Multi-GPU index:
device = 'cuda:0'        # first GPU
n_gpus = torch.cuda.device_count()

# ===== Create directly on device =====
x = torch.randn(3, 3, device=device)
# Faster than creating on CPU then moving.

# ===== Move tensors =====
x_cpu = torch.zeros(3, 3)
x_gpu = x_cpu.to(device)
x_gpu = x_cpu.to(device, non_blocking=True)
x_gpu = x_cpu.cuda()                     # legacy; .to is preferred

# ===== Move modules =====
model = torch.nn.Linear(4, 3).to(device)

# Every Variable / Parameter must be on the SAME device as input tensors.
# Common error message: 'expected all tensors to be on the same device'.

# ===== DataLoader + pinned memory =====
from torch.utils.data import DataLoader, TensorDataset
ds = TensorDataset(torch.randn(100, 4), torch.randint(0, 3, (100,)))
loader = DataLoader(ds, batch_size=8, pin_memory=True, num_workers=2)

for X, y in loader:
    X = X.to(device, non_blocking=True)
    y = y.to(device, non_blocking=True)
    # train ...

# pin_memory=True + non_blocking=True together let the GPU copy in parallel with compute.

# ===== Mixed precision (autocast) =====
from torch.amp import autocast, GradScaler
scaler = GradScaler(device='cuda')

for X, y in loader:
    X = X.to(device, non_blocking=True); y = y.to(device, non_blocking=True)
    optim.zero_grad(set_to_none=True)
    with autocast(device_type='cuda', dtype=torch.float16):
        out = model(X)
        loss = loss_fn(out, y)
    scaler.scale(loss).backward()
    scaler.step(optim)
    scaler.update()

# ===== Apple Silicon (MPS) =====
# Same patterns; device = 'mps'.
# Not every op is implemented — fall back to CPU may be needed for some models.
# PyTorch sets PYTORCH_ENABLE_MPS_FALLBACK=1 to auto-fall-back.

# ===== Querying GPU info =====
torch.cuda.is_available()
torch.cuda.device_count()
torch.cuda.get_device_name(0)
torch.cuda.memory_allocated()              # bytes
torch.cuda.max_memory_allocated()
torch.cuda.empty_cache()                    # release cached allocator memory

# ===== DistributedDataParallel (multi-GPU) =====
# Use torch.distributed.init_process_group + DDP wrapper.
# Run via torchrun:
# torchrun --nproc_per_node=4 train.py

# ===== Patterns to internalise =====
# - Detect device once at the top of main; thread it through
# - .to(device, non_blocking=True) + pin_memory=True is the perf trinity
# - autocast + GradScaler for mixed precision
# - Set seeds AND torch.backends.cudnn.deterministic for repro

# ===== Pitfalls =====
# - Model on GPU, input on CPU -> device mismatch error
# - Forgetting non_blocking + pin_memory -> CPU/GPU stalls
# - GradScaler with fp32 model -> no help (fp16 only)
# - torch.cuda calls without CUDA -> errors on Mac / pure CPU

import torch
dev = 'cuda' if torch.cuda.is_available() else ('mps' if torch.backends.mps.is_available() else 'cpu')
x = torch.randn(1024, 1024, device=dev)