58. PyTorch: tensors, nn.Module, and the training loop

58.1 Tensors and devices

import torch

device = (
    "cuda" if torch.cuda.is_available()          # NVIDIA GPU
    else "mps" if torch.backends.mps.is_available()   # Apple Silicon GPU
    else "cpu"
)
x = torch.randn(4, 3, device=device)             # like np.random.normal, but can live on a GPU
w = torch.randn(3, 2, device=device, requires_grad=True)   # track gradients for w
y = (x @ w).relu().sum()
y.backward()
print(device, y.item(), w.grad.shape)            # .item() copies a scalar back to Python
print(torch.from_numpy(x.cpu().numpy()).shape)   # NumPy <-> torch interop (CPU tensors share memory)

A tensor is an ndarray that can live on a GPU and that records the operations applied to it, exactly like Value in Ch 57. Moving data between CPU and GPU (.to(device), .cpu(), .item()) is slow relative to computing; keep data on the device for the whole loop, and avoid .item() inside the hot path.

58.2 A complete training loop

# labs/train_mlp.py: train a small classifier the standard PyTorch way
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset, random_split

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

# Data: points inside/outside a circle (the tinygrad task, now with 10 000 samples)
X = torch.rand(10_000, 2) * 2 - 1
y = ((X ** 2).sum(dim=1) < 0.5).long()               # class 0 or 1
train_ds, val_ds = random_split(TensorDataset(X, y), [8000, 2000])
train_dl = DataLoader(train_ds, batch_size=128, shuffle=True)
val_dl = DataLoader(val_ds, batch_size=512)

class MLP(nn.Module):
    def __init__(self, hidden: int = 64) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
            nn.Linear(hidden, 2),                      # logits for 2 classes
        )
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

model = MLP().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=1e-4)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(10):
    model.train()                                      # enables dropout/batch-norm training behaviour
    for xb, yb in train_dl:
        xb, yb = xb.to(device), yb.to(device)
        logits = model(xb)                             # 1. forward
        loss = loss_fn(logits, yb)                     # 2. loss
        opt.zero_grad(set_to_none=True)                # 3. clear old gradients
        loss.backward()                                # 4. backprop
        opt.step()                                     # 5. update

    model.eval()
    correct = total = 0
    with torch.no_grad():                              # no graph building during evaluation
        for xb, yb in val_dl:
            pred = model(xb.to(device)).argmax(dim=1)
            correct += (pred == yb.to(device)).sum().item()
            total += len(yb)
    print(f"epoch {epoch}  train loss {loss.item():.4f}  val acc {correct / total:.3f}")

torch.save(model.state_dict(), "mlp.pt")              # weights only; see Ch 47.1 for loading safely

Every PyTorch training script, up to the LLM you train in AI Ch 41, has this shape: forward, loss, zero_grad, backward, step, plus an evaluation pass under torch.no_grad() and model.eval(). What changes for bigger models is around the loop, not the loop itself:

58.3 Saving models safely

from safetensors.torch import save_file, load_file   # uv add safetensors

save_file(model.state_dict(), "mlp.safetensors")     # tensors only: no code, mmap-able
state = load_file("mlp.safetensors")
model.load_state_dict(state)
# torch.load("third_party.pt")  -> defaults to weights_only=True since PyTorch 2.6; keep it that way