PyTorch Mental Model

The four PyTorch primitives you actually use

Tensor — n-dimensional array on a device, with a dtype, shape, and optional gradient tracking. The atomic unit.
Module — a stateful piece of the network that holds parameters and defines a forward pass. Composed via attributes (self.conv = nn.Conv2d(...)) and discovered automatically by state_dict() and parameters().
Optimizer — knows about a list of parameters and updates them given their gradients.
Autograd engine — builds a computation graph during forward, walks it backward to compute gradients.

Every PyTorch program you read is one of: defining a Module, putting tensors through it, computing a scalar loss, calling .backward(), calling opt.step(). Once you internalize that loop, the rest is library knowledge.

Tip: PyTorch is a thin layer of conventions over numpy + automatic differentiation. The framework is small, the ecosystem (torchvision, transformers, lightning) is huge. Learn the core deeply; pick the ecosystem pieces by need.

Modules compose like Lego

A Module can contain other Modules. state_dict() recursively flattens all parameters, to(device) recursively moves them. The hierarchy is the API — you don't have to register parameters manually as long as you assign them as attributes.

The 2026 torch.compile detail

Wrap any Module in torch.compile(model) and PyTorch traces the forward pass into a fused graph that runs significantly faster on modern accelerators. First call is slower (compilation); subsequent calls are 1.3-3x faster. Pure win on CUDA; smaller win on CPU/MPS.

Principle: PyTorch is the lingua franca of deep learning in 2026. Reading and writing fluent PyTorch is the single highest-leverage skill in this quest.

Code

Modules compose, parameters are auto-discovered·python

import torch
import torch.nn as nn

class Block(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.fc = nn.Linear(dim, dim)
        self.norm = nn.LayerNorm(dim)
    def forward(self, x):
        return self.norm(torch.relu(self.fc(x)))

class Net(nn.Module):
    def __init__(self, dim, depth):
        super().__init__()
        self.blocks = nn.ModuleList([Block(dim) for _ in range(depth)])
        self.head = nn.Linear(dim, 10)
    def forward(self, x):
        for blk in self.blocks:
            x = x + blk(x)            # residual
        return self.head(x)

m = Net(dim=128, depth=4)
print(sum(p.numel() for p in m.parameters()))
m.to("cuda" if torch.cuda.is_available() else "cpu")

torch.compile for free 1.3-3x speedup·python

import torch
model = MyModel().to("cuda")
compiled = torch.compile(model, mode="reduce-overhead")
# Use compiled exactly like model — same API
out = compiled(x)

The four PyTorch primitives you actually use

Modules compose like Lego

The 2026 torch.compile detail

Code

External links

Exercise

Progress

Comments 0