PyTorch Basics: Tensors, Autograd, nn.Module

PyTorch is the most-used DL framework in research. 90% of papers ship PyTorch code.

If you learn NumPy → PyTorch is “NumPy + autograd + GPU”.

This piece: the essentials to read and write PyTorch.

Tensors

PyTorch’s core data structure. Like NumPy arrays:

import torch

# Create
x = torch.tensor([1.0, 2.0, 3.0])
y = torch.zeros(3, 4)
z = torch.ones(2, 2)
r = torch.randn(3, 3)              # standard normal
ar = torch.arange(0, 10, 2)        # [0, 2, 4, 6, 8]

# From NumPy
import numpy as np
arr = np.array([1, 2, 3])
t = torch.from_numpy(arr)

# Back to NumPy
arr = t.numpy()

Tensor Operations

Just like NumPy:

a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([10.0, 20.0, 30.0])

a + b      # [11, 22, 33]
a * b      # element-wise
a @ b      # dot product / matmul
a.sum()
a.mean()
a.reshape(3, 1)
a.unsqueeze(0)   # add dim: [1, 2, 3] → [[1, 2, 3]]

GPU Support

The killer feature:

device = "cuda" if torch.cuda.is_available() else "cpu"
# Or for Apple Silicon: "mps"

x = torch.randn(1000, 1000, device=device)
y = torch.randn(1000, 1000, device=device)
z = x @ y    # runs on GPU, 100× faster than CPU

Move existing tensors:

x = x.to(device)
model = model.to(device)

Rule: model + all input tensors must be on the same device.

Autograd: Automatic Differentiation

The real magic.

x = torch.tensor(2.0, requires_grad=True)

y = x**2 + 3*x + 1   # y = 11 at x=2

y.backward()         # compute gradients
print(x.grad)        # dy/dx = 2x + 3 = 7 ✓

requires_grad=True marks the tensor as “track operations”. .backward() computes gradients.

This is L1-06 / L3-02 in action — backprop, automated.

Neural Network Modules

PyTorch’s nn module contains layer building blocks:

import torch.nn as nn

# Single layer
linear = nn.Linear(in_features=100, out_features=10)

# Common layers
nn.ReLU()
nn.Dropout(p=0.5)
nn.BatchNorm1d(64)
nn.Conv2d(3, 32, kernel_size=3)
nn.LSTM(input_size=100, hidden_size=50)
nn.MultiheadAttention(embed_dim=512, num_heads=8)

Building a Model

Subclass nn.Module:

class MyMLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)

# Instantiate
model = MyMLP(input_dim=784, hidden_dim=256, output_dim=10)
print(model)

forward defines the computation. PyTorch handles backward automatically.

Loss Functions

nn.CrossEntropyLoss()   # multi-class classification
nn.BCEWithLogitsLoss()  # binary classification
nn.MSELoss()            # regression
nn.L1Loss()             # L1 regression
nn.SmoothL1Loss()       # Huber
nn.KLDivLoss()          # distribution matching

CrossEntropyLoss is the most-used.

Optimizers

import torch.optim as optim

optimizer = optim.Adam(model.parameters(), lr=0.001)
# or:
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)

Adam for most cases. AdamW for transformers.

The Training Loop

The pattern you’ll use 1000 times:

for epoch in range(num_epochs):
    model.train()
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)

        # Forward
        pred = model(x)
        loss = loss_fn(pred, y)

        # Backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # Eval
    model.eval()
    with torch.no_grad():
        for x, y in val_loader:
            x, y = x.to(device), y.to(device)
            pred = model(x)
            # compute val loss / accuracy

Key elements:

• model.train() / model.eval() — toggles dropout, batchnorm
• optimizer.zero_grad() — clear previous gradients
• loss.backward() — compute new gradients
• optimizer.step() — apply gradient update
• torch.no_grad() — disable autograd for inference (saves memory)

DataLoader

PyTorch’s data pipeline:

from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

dataset = MyDataset(X, y)
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

num_workers > 0 parallelizes data loading on CPU while GPU trains.

Built-in Datasets

from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_data = datasets.MNIST(".", train=True, download=True, transform=transform)
loader = DataLoader(train_data, batch_size=64, shuffle=True)

Convenient for benchmarks and tutorials.

Saving / Loading Models

# Save weights
torch.save(model.state_dict(), "model.pt")

# Load weights (rebuild model first)
model = MyMLP(...)
model.load_state_dict(torch.load("model.pt"))
model.eval()

# Save entire model (less portable)
torch.save(model, "full_model.pt")

Standard practice: save state_dict (weights only), not the full pickled object.

Inference Mode

For production:

model.eval()
with torch.inference_mode():   # newer than no_grad, slightly faster
    pred = model(x)

A Complete Example: MNIST Classifier

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

device = "cuda" if torch.cuda.is_available() else "cpu"

# Data
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train = datasets.MNIST(".", train=True, download=True, transform=transform)
test = datasets.MNIST(".", train=False, transform=transform)
train_loader = DataLoader(train, batch_size=128, shuffle=True)
test_loader = DataLoader(test, batch_size=128)

# Model
class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        )
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128), nn.ReLU(),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        return self.fc(self.conv(x))

model = CNN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# Train
for epoch in range(5):
    model.train()
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        logits = model(x)
        loss = criterion(logits, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # Eval
    model.eval()
    correct = 0
    with torch.no_grad():
        for x, y in test_loader:
            x, y = x.to(device), y.to(device)
            pred = model(x).argmax(dim=1)
            correct += (pred == y).sum().item()
    print(f"Epoch {epoch}: test acc = {correct / len(test):.3f}")

After 5 epochs: ~99% test accuracy. ~3 min on a modern GPU.

This is how modern DL works in 50 lines.

Common Bugs

1. Forgot zero_grad

# BAD: gradients accumulate
for ...:
    loss.backward()
    optimizer.step()   # uses stale + new gradients

# GOOD
for ...:
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

2. Device mismatch

x = torch.randn(10, device="cuda")
y = torch.randn(10)                # on CPU
z = x + y   # ERROR

Keep everything on the same device.

3. Forgot eval mode

# Test phase
model.eval()    # ← critical, disables dropout / batchnorm updates
with torch.no_grad():
    ...

Without eval(), dropout still drops neurons during inference — bad results.

4. Tensor vs scalar in loss

loss = (pred - target) ** 2     # tensor of losses
loss.backward()                  # ERROR: must be scalar
loss.mean().backward()           # correct

PyTorch Ecosystem

• torchvision — image datasets, models, transforms
• torchaudio — audio
• torchtext — NLP (somewhat legacy now)
• transformers (HuggingFace) — pretrained LLMs / vision models
• lightning — high-level training framework
• fastai — opinionated wrapper, great for learners

For most projects: PyTorch + transformers + lightning is a strong stack.

Compile (PyTorch 2.0+)

For production speedup:

model = torch.compile(model)   # graph-based optimization, 30-100% faster

Adds ~1min compile time on first run, then much faster.

Next recommended: L3-01 Perceptron to MLP or L1-08 NumPy.