Introduction to PyTorch

By the end of this lesson, you will be able to...

• trace one PyTorch training batch from tensors through loss, gradients, optimizer state, and updated parameters;

• distinguish train()/eval() mode from gradient tracking and explain why both matter;

• diagnose a basic training-loop error using shapes and state ownership instead of treating framework calls as magic.

Idea in one sentence: PyTorch automates tensor bookkeeping and reverse-mode differentiation, but a reliable training loop still has the same visible owners and state transitions as the manual network you already traced.

Core Insight

Consider a small binary classifier whose validation score will not improve. The model definition looks ordinary, and loss.backward() runs without an error. The trouble is in the loop: old gradients were never cleared, validation ran with dropout still active, and the target tensor has shape [32] while the model returns logits of shape [32, 1].

It is tempting to see PyTorch as a single machine that “trains the model.” That model works for a quick demo, but it fails the moment a result is surprising. The framework has not removed the training mechanism. It has divided it into objects that own different kinds of state.

The stronger model is a state ledger: the module owns parameters and mode, autograd owns the graph for the current forward pass, parameters hold accumulated gradients, and the optimizer owns update state. A training loop is the controlled handoff between them.

The Small Model and Its Owners

We will use a tiny classifier with two input features and one output logit. A logit is an unnormalized binary score; BCEWithLogitsLoss combines the sigmoid interpretation and binary cross-entropy in one numerically stable loss.

import torch
from torch import nn

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(2, 4)
        self.relu = nn.ReLU()
        self.out = nn.Linear(4, 1)

    def forward(self, x):
        return self.out(self.relu(self.hidden(x)))

model = TinyNet()
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

Assigning nn.Linear modules to self makes them children of TinyNet. Their weights and biases become registered model parameters, so model.parameters() gives the optimizer the tensors it is allowed to update. forward() describes the computation, but it does not perform an update by itself.

For one illustrative batch, use these shapes:

x_batch: [32, 2]  float features
y_batch: [32, 1]  float targets containing 0 or 1
logits:  [32, 1]  one raw score per example

Shape is part of the contract. Here the target has the same shape as the logits because this loss compares one binary score with one binary target for every example. A target shaped [32] may be a sign that a dimension was dropped; fix the data contract deliberately rather than hoping broadcasting means what the task means.

The Initial Model: backward() Trains the Network

The initial mental model is understandable: call the model, compute a loss, call backward(), and the network has learned. That is close enough to observe a loss value, but it misses two state changes.

First, backward() computes derivatives and accumulates them in the .grad fields of leaf parameters. It does not change the parameter values. Second, an optimizer such as Adam has its own moving-average state in addition to the parameter tensors. optimizer.step() reads gradients and changes parameters according to that optimizer's rule.

If the next batch calls backward() without clearing existing gradients, new gradients add to the old ones. That accumulation is useful for deliberate gradient accumulation across several micro-batches, but it is a bug when the intended update is one batch at a time.

The missing model is not another API call. It is ownership: one object computes, another stores derivatives, and another applies updates.

One Batch, Step by Step

Here is a small complete training step. The values are not a recommendation about model size or learning rate; they make the transitions visible.

model.train()

optimizer.zero_grad(set_to_none=True)
logits = model(x_batch)
loss = loss_fn(logits, y_batch)
loss.backward()
optimizer.step()

Read it as a trace, not as an incantation.

The order matters. Calling optimizer.step() before backward() gives it no newly computed gradient to use. Clearing gradients after backward() but before step() throws away the evidence of this batch. Calling zero_grad() at the start makes the intended ownership easy to inspect.

Ask one useful question at the middle of the trace: where did the graph go after step()? In ordinary PyTorch training, the forward graph is used to compute gradients and then released. The persistent state is the model parameters, their current gradients until cleared, and the optimizer state—not a complete history of every batch.

So far, we have seen that framework code is compact because it delegates work to objects with specific state. This matters because a strange result can be traced by asking which object owned the state at the moment it changed.

When a Gradient Is Missing

After loss.backward(), inspect a parameter rather than trusting that the call succeeded:

print(model.hidden.weight.grad is None)

For a parameter that influenced the current loss, the expected answer is normally False: it has a gradient tensor. None has a more specific meaning than “the gradient happened to be zero.” A zero gradient is a computed value. None usually means this parameter was not part of the differentiable path from the current loss, its gradient was disabled, or gradients were cleared before inspection.

Use a small ownership trace to narrow the cause:

parameter -> forward operation -> logits -> loss -> backward -> parameter.grad

If forward() bypasses self.hidden, its weights have no path to the loss and receive no gradient. If code converts a tensor to a Python number or detaches it before computing the loss, that path no longer carries autograd history. If a layer is intentionally frozen, its parameters may have requires_grad=False, which is correct only when the freezing decision is explicit.

Do not expect every intermediate activation to retain a convenient .grad field after backward. PyTorch's usual interface is optimized around gradients of leaf tensors such as model parameters. When debugging, first inspect the loss, then the named parameter gradients, then the input/output shapes around the suspicious layer. That sequence distinguishes a disconnected computation from a legitimate but small update without turning gradient inspection into guesswork.

Training Mode Is Not Gradient Mode

Two switches are often confused because both appear around validation and inference:

model.eval()
with torch.no_grad():
    logits = model(x_val)
    probabilities = torch.sigmoid(logits)

model.eval() changes the behavior of modules whose forward pass depends on mode. In this track, the important examples are Dropout, which stops sampling masks, and BatchNorm, which uses its stored running statistics instead of current batch statistics. It does not turn gradients off.

torch.no_grad() disables recording operations for autograd inside its block. It reduces unnecessary graph and gradient work for a standard evaluation forward pass. It does not set Dropout or BatchNorm to evaluation behavior.

The correct validation path usually uses both. There are exceptions—such as deliberately computing input gradients—but that is a named requirement, not a reason to omit either switch by accident.

Make the Framework Map Explicit

The manual lessons did not become obsolete when the code got shorter. They became a debugging map.

For example, suppose validation loss is erratic. Before changing the architecture, print model.training, confirm that validation has a deterministic preprocessing path, and verify whether each validation batch sees eval() plus no_grad(). Those checks connect the previous lessons on BatchNorm, Dropout, and augmentation to the framework rather than treating them as isolated topics.

Where This Compactness Can Mislead You

PyTorch makes it easy to run many experiments, but it does not choose a sensible loss, preserve label semantics, create a representative validation split, or interpret an improving metric. Those are model and data decisions.

The trade-off is speed versus invisible assumptions. Registered modules, automatic gradients, and optimizer objects reduce handwritten code; they also make it possible to execute an incorrect training loop quickly. A framework error is often an ordinary learning-system error in a shorter costume: wrong shapes, wrong loss-target pairing, stale gradients, incorrect mode, or an update applied to the wrong parameters.

You can see the boundary when the code runs but the evidence is incoherent: loss decreases while validation collapses, identical inputs yield varying validation predictions, a parameter's gradient is None, or optimizer steps leave parameters unchanged. Inspect state and contracts first; do not assume a successful Python execution proves a valid learning procedure.

Checks

Check: What is wrong with this order?

loss = loss_fn(model(x_batch), y_batch)
loss.backward()
optimizer.zero_grad()
optimizer.step()

Think first, then reveal.

Answer: zero_grad() clears the gradients just computed by backward() before step() can use them. Clear at the start of the next intended update, then compute loss, call backward(), and call step().

Check: A validation loop uses torch.no_grad() but leaves model.train() active. What can still vary unexpectedly?

Answer: Dropout can still sample masks and BatchNorm can still use training behavior. no_grad() controls graph recording, not module mode; call model.eval() too.

Check: The model returns logits shaped [32, 1], but targets are [32]. What should you inspect before reshaping something?

Answer: Check the task's intended output unit and the loss contract. For one binary target per example, make both tensors explicitly [32, 1] (or choose a consistent one-dimensional design); do not rely on an accidental broadcast.

Transfer: Repair a Validation Step

This validation code runs without raising an exception, but predictions vary between two calls on the same batch:

for x_val, y_val in val_loader:
    logits = model(x_val)
    val_loss += loss_fn(logits, y_val)

Write the smallest repair and explain which state each added line controls.

Model answer: Before the loop, call model.eval() to switch Dropout and BatchNorm to evaluation behavior. Wrap the loop in with torch.no_grad(): to stop recording graphs that validation does not need. Keep the validation preprocessing deterministic and ensure y_val has the loss-compatible shape. eval() controls module behavior; no_grad() controls autograd recording. Neither line changes weights, so resume model.train() before the next training epoch.

Resources

• [DOCUMENTATION] PyTorch: Modules — Focus: registered submodules, parameters, and train/eval mode.
• [DOCUMENTATION] PyTorch: Autograd Mechanics — Focus: graph recording and backward differentiation.
• [TUTORIAL] PyTorch: Optimizing Model Parameters — Focus: zero_grad, backward, and step in one training batch.
• [TUTORIAL] PyTorch: Saving and Loading Models — Focus: model and optimizer state boundaries that the next lesson will organize.

Key Takeaways

• A PyTorch training batch still follows forward → loss → backward → update, with gradients and optimizer state owned separately from model parameters.
• model.eval() controls mode-dependent layers, while torch.no_grad() controls autograd recording; standard validation needs both.
• Treat tensor shapes, gradients, modes, and optimizer state as inspectable contracts. Compact code does not make those responsibilities disappear.