Install / CUDA

# ===== 1. Python env =====
python -m venv .venv
source .venv/bin/activate
pip install --upgrade pip

# ===== 2. Install (CPU, all platforms) =====
pip install torch torchvision torchaudio

# ===== 3. macOS (Apple Silicon: MPS backend) =====
pip install torch torchvision torchaudio
# MPS is the GPU on M-series Macs. PyTorch ships ready.

# Verify MPS:
python - <<'PY'
import torch
print('cuda:', torch.cuda.is_available())
print('mps:', torch.backends.mps.is_available())
PY

# ===== 4. Linux / Windows with NVIDIA GPU =====
# Check pytorch.org for the latest selector — versions move fast.
# Example (CUDA 12.1):
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# ===== 5. ROCm (AMD GPU on Linux) =====
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.7

# ===== 6. Docker (the most painless GPU path) =====
docker run --gpus all -it --rm pytorch/pytorch:2.3-cuda12-cudnn8-runtime bash

# ===== 7. Verify =====
python - <<'PY'
import torch
print('torch', torch.__version__)
print('cuda', torch.cuda.is_available())
print('mps', torch.backends.mps.is_available())

device = 'cuda' if torch.cuda.is_available() else ('mps' if torch.backends.mps.is_available() else 'cpu')
print('device:', device)
x = torch.randn(3, 3).to(device)
print(x @ x.T)
PY

# ===== 8. Smoke train =====
python - <<'PY'
import torch, torch.nn as nn
device = 'cuda' if torch.cuda.is_available() else ('mps' if torch.backends.mps.is_available() else 'cpu')
model = nn.Sequential(nn.Linear(4, 32), nn.ReLU(), nn.Linear(32, 3)).to(device)
optim = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
X = torch.randn(100, 4).to(device); y = torch.randint(0, 3, (100,)).to(device)
for _ in range(5):
    optim.zero_grad()
    loss = loss_fn(model(X), y)
    loss.backward()
    optim.step()
print('final loss:', loss.item())
PY

# ===== 9. Common companions =====
pip install transformers accelerate datasets evaluate
pip install pytorch-lightning
pip install tensorboard wandb

# ===== Patterns to internalise =====
# - Always check pytorch.org for the matching CUDA index URL
# - Smoke-train a tiny model after install; catches device issues early
# - Pin in requirements.txt with the index URL noted in a comment
# - For LLM work, Hugging Face transformers + accelerate are the floor

# ===== Pitfalls =====
# - Installing torch without the CUDA index URL -> CPU-only wheels
# - Mixing torch + torchvision versions; they must match major
# - Using cudatoolkit from conda and pip wheels together -> conflicts
# - GPU detected but slow -> wrong device strings (use the device variable consistently)

# CPU
pip install torch torchvision torchaudio
# CUDA — see pytorch.org/get-started for your platform
# Verify
python -c 'import torch; print(torch.__version__, torch.cuda.is_available())'