PyTorch
Keras
training loop
machine learning
deep learning

PyTorch is there a definitive training loop similar to Keras' fit?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

PyTorch does not have one single low-level training loop that everyone must use, but it does have a very standard pattern. If you are coming from Keras, the closest answer is: there is no single built-in fit() style loop at the same abstraction level, yet there is a widely accepted manual loop shape that most PyTorch projects follow.

The Canonical PyTorch Training Loop

A typical PyTorch training loop has the same moving parts as Keras fit():

  • iterate over batches from a DataLoader
  • run the forward pass
  • compute the loss
  • zero the gradients
  • call backward()
  • step the optimizer

Here is a compact example for classification:

python
1import torch
2from torch import nn
3from torch.utils.data import DataLoader, TensorDataset
4
5x = torch.randn(100, 10)
6y = torch.randint(0, 3, (100,))
7
8dataset = TensorDataset(x, y)
9loader = DataLoader(dataset, batch_size=16, shuffle=True)
10
11model = nn.Sequential(
12    nn.Linear(10, 32),
13    nn.ReLU(),
14    nn.Linear(32, 3),
15)
16
17loss_fn = nn.CrossEntropyLoss()
18optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
19
20for epoch in range(5):
21    model.train()
22    epoch_loss = 0.0
23
24    for xb, yb in loader:
25        optimizer.zero_grad()
26        logits = model(xb)
27        loss = loss_fn(logits, yb)
28        loss.backward()
29        optimizer.step()
30        epoch_loss += loss.item()
31
32    print(f"epoch={epoch} loss={epoch_loss / len(loader):.4f}")

This is the pattern most PyTorch users mean when they talk about "the training loop."

Why PyTorch Leaves It Explicit

PyTorch has historically favored flexibility over a one-size-fits-all training entry point. That makes unusual research loops easier to express:

  • multiple optimizers
  • gradient accumulation
  • custom mixed precision logic
  • reinforcement learning updates
  • alternating generator and discriminator steps

In Keras, fit() is excellent when your problem matches the standard supervised-learning pattern. In PyTorch, the framework gives you the building blocks and expects you to assemble the loop that matches your problem.

What Counts as "Definitive"

There is no official one true loop because the right loop depends on the workload. But there is a canonical baseline:

  1. call model.train() for training
  2. iterate over the loader
  3. zero gradients before each backward pass
  4. compute predictions and loss
  5. call loss.backward()
  6. call optimizer.step()

That sequence is stable across a huge portion of PyTorch codebases.

For validation, the pattern changes slightly:

python
1model.eval()
2
3with torch.no_grad():
4    for xb, yb in loader:
5        logits = model(xb)
6        loss = loss_fn(logits, yb)

model.eval() changes layer behavior for modules such as dropout and batch normalization, while torch.no_grad() avoids building gradient graphs during evaluation.

If You Want a fit() Experience

If your goal is not maximum loop control, higher-level wrappers exist. Libraries such as PyTorch Lightning and similar orchestration frameworks provide a more Keras-like experience on top of PyTorch. They are useful when you want callbacks, checkpoints, and logging without rewriting the same loop structure in every project.

The tradeoff is abstraction. The more you hide, the less obvious custom training behavior becomes.

A Good Practical Strategy

For learning and small projects, write the loop manually at least once. That teaches you what Keras fit() is normally doing for you.

After that, choose based on project needs:

  • manual loop for custom or research-heavy training
  • high-level wrapper for productivity and standard workflows

The important part is understanding the moving pieces rather than memorizing one exact template.

Common Pitfalls

The most common mistake is forgetting optimizer.zero_grad(), which causes gradients to accumulate across batches unintentionally.

Another issue is skipping model.eval() during validation. That can make metrics look inconsistent because dropout and batch normalization behave differently in training mode.

Developers also expect PyTorch to have one official fit() equivalent and then feel blocked by the absence of one. In practice, the standard loop is short, readable, and often a benefit rather than a burden.

Summary

  • PyTorch has a canonical training-loop pattern, but not one single built-in low-level fit() equivalent like Keras.
  • The standard loop is forward pass, loss, backward pass, and optimizer step over batches.
  • Validation uses model.eval() and torch.no_grad().
  • PyTorch keeps the loop explicit so custom training behavior stays easy to express.
  • If you want a higher-level fit() style workflow, use a wrapper library on top of PyTorch.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.