Why the training loop is explicit in PyTorch

Introduction

PyTorch makes you write the training loop yourself so you can see and control every step clearly. This helps you learn and customize your model easily.

When you want to understand how your model learns step-by-step.

When you need to customize training for special tasks or data.

When debugging your model to find and fix problems.

When experimenting with new ideas that need control over training.

When learning machine learning basics and want to see what happens inside.

Syntax

PyTorch

for epoch in range(num_epochs):
    for inputs, labels in dataloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()

This loop runs over your data many times (epochs).

You calculate loss, then update model weights manually.

Examples

Basic training loop running 3 times over data.

PyTorch

for epoch in range(3):
    for x, y in train_loader:
        optimizer.zero_grad()
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        loss.backward()
        optimizer.step()

Same loop with different variable names and 5 epochs.

PyTorch

for epoch in range(5):
    for inputs, targets in data_loader:
        optimizer.zero_grad()
        predictions = model(inputs)
        loss = criterion(predictions, targets)
        loss.backward()
        optimizer.step()

Sample Model

This code trains a simple model to learn y = 2x + 1. The training loop is explicit so you see each step: zeroing gradients, forward pass, loss calculation, backward pass, and optimizer step. After training, it predicts the output for input 5.0.

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# Simple dataset: y = 2x + 1
x_train = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y_train = torch.tensor([[3.0], [5.0], [7.0], [9.0]])

dataset = TensorDataset(x_train, y_train)
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

# Simple linear model
model = nn.Linear(1, 1)

# Loss and optimizer
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

num_epochs = 10

for epoch in range(num_epochs):
    total_loss = 0
    for inputs, labels in dataloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_fn(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {total_loss:.4f}")

# Test prediction
test_input = torch.tensor([[5.0]])
prediction = model(test_input).item()
print(f"Prediction for input 5.0: {prediction:.2f}")

Important Notes

Explicit loops help you understand what happens inside training.

You can change any step easily, like adding extra calculations.

Other libraries hide this loop, but PyTorch shows it for learning and flexibility.

Summary

PyTorch uses explicit training loops so you control every step.

This helps learning, debugging, and customizing models.

Writing the loop yourself shows how training really works.

Practice

1. Why does PyTorch require you to write the training loop explicitly instead of hiding it inside a function?

easy

A. Because it gives you full control over each step of training.

B. Because PyTorch does not support loops.

C. Because training loops are automatically handled by PyTorch.

D. Because explicit loops make the code run faster by default.

Solution

1. Step 1: Understand PyTorch's design philosophy

   PyTorch is designed to be flexible and transparent, letting users control every detail of the training process.

2. Step 2: Compare with other frameworks

   Unlike some frameworks that hide training loops, PyTorch requires explicit loops to allow customization like dynamic graphs and debugging.

3. Final Answer:

   Because it gives you full control over each step of training. -> Option A

4. Quick Check:

   Explicit loop = full control [OK]

Hint: Explicit loops mean you control every training step [OK]

Common Mistakes:

• Thinking PyTorch hides the training loop automatically
• Believing loops are not supported in PyTorch
• Assuming explicit loops always make code faster

2. Which of the following is the correct way to start a training loop in PyTorch?

easy

A. train_loop(num_epochs)

B. while epoch < num_epochs:

C. for epoch in range(num_epochs):

D. loop over epochs automatically

Solution

1. Step 1: Identify the common Python loop syntax

   PyTorch training loops typically use Python's for loop with range to iterate over epochs.

2. Step 2: Check each option

   for epoch in range(num_epochs): uses the correct for loop syntax; B is valid Python but less common; C and D are not valid PyTorch syntax.

3. Final Answer:

   for epoch in range(num_epochs): -> Option C

4. Quick Check:

   Use for loop with range for epochs [OK]

Hint: Use for loop with range to iterate epochs [OK]

Common Mistakes:

• Trying to call a non-existent train_loop function
• Expecting PyTorch to loop automatically
• Using while loops without proper epoch increment

3. Given this PyTorch training loop snippet, what will be printed?

for epoch in range(2):
    print(f"Epoch {epoch+1} start")
    for batch in range(3):
        print(f" Batch {batch+1} processed")

medium

A. Epoch 1 start\n Batch 1 processed\n Batch 2 processed\n Batch 3 processed

B. Epoch 0 start\n Batch 0 processed\n Batch 1 processed\n Batch 2 processed\nEpoch 1 start\n Batch 0 processed\n Batch 1 processed\n Batch 2 processed

C. Epoch 1 start\n Batch 1 processed\n Batch 2 processed\nEpoch 2 start\n Batch 1 processed\n Batch 2 processed

D. Epoch 1 start\n Batch 1 processed\n Batch 2 processed\n Batch 3 processed\nEpoch 2 start\n Batch 1 processed\n Batch 2 processed\n Batch 3 processed

Solution

1. Step 1: Analyze the outer loop

   The outer loop runs for 2 epochs: epoch values 0 and 1, printing "Epoch {epoch+1} start" which gives 1 and 2.

2. Step 2: Analyze the inner loop

   For each epoch, the inner loop runs 3 times (batch 0 to 2), printing "Batch {batch+1} processed" which gives batches 1, 2, 3.

3. Final Answer:

   Epoch 1 start Batch 1 processed Batch 2 processed Batch 3 processed Epoch 2 start Batch 1 processed Batch 2 processed Batch 3 processed -> Option D

4. Quick Check:

   2 epochs x 3 batches each = 6 batch prints [OK]

Hint: Count loops and add 1 to zero-based indices [OK]

Common Mistakes:

• Forgetting to add 1 to epoch or batch for printing
• Confusing zero-based indexing with one-based output
• Missing the second epoch's batch prints

4. What is wrong with this PyTorch training loop snippet?

for epoch in range(3):
    optimizer.zero_grad()
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward
    optimizer.step()

medium

A. optimizer.zero_grad() should be called after optimizer.step().

B. The loss.backward call is missing parentheses, so gradients are not computed.

C. The loop should use while instead of for.

D. loss_fn should be called after optimizer.step().

Solution

1. Step 1: Check the loss.backward call

   loss.backward is a function and must be called with parentheses: loss.backward(). Without parentheses, gradients won't be computed.

2. Step 2: Verify other statements

   optimizer.zero_grad() is correctly placed before backward; optimizer.step() is correctly after backward; loop type is fine; loss_fn call order is correct.

3. Final Answer:

   The loss.backward call is missing parentheses, so gradients are not computed. -> Option B

4. Quick Check:

   Call loss.backward() with parentheses [OK]

Hint: Remember to call loss.backward() with () [OK]

Common Mistakes:

• Forgetting parentheses on loss.backward
• Calling zero_grad after optimizer.step
• Changing loop type unnecessarily

5. You want to customize your PyTorch training loop to print the average loss every 5 batches. Which code snippet correctly implements this inside the explicit training loop?

hard

A. for epoch in range(num_epochs):
running_loss = 0.0
for i, data in enumerate(dataloader, 1):
inputs, labels = data
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_fn(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 5 == 0:
print(f'Batch {i}, Avg Loss: {running_loss / 5:.4f}')
running_loss = 0.0

B. for epoch in range(num_epochs):
running_loss = 0.0
for i, data in enumerate(dataloader):
inputs, labels = data
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_fn(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 5 == 0:
print(f'Batch {i}, Avg Loss: {running_loss / 5:.4f}')
running_loss = 0.0

C. for epoch in range(num_epochs):
running_loss = 0.0
for i, data in enumerate(dataloader, 1):
inputs, labels = data
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_fn(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 4 == 0:
print(f'Batch {i}, Avg Loss: {running_loss / 5:.4f}')
running_loss = 0.0

D. for epoch in range(num_epochs):
running_loss = 0.0
for i, data in enumerate(dataloader, 1):
inputs, labels = data
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_fn(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 5 == 1:
print(f'Batch {i}, Avg Loss: {running_loss / 5:.4f}')
running_loss = 0.0

Solution

1. Step 1: Check enumerate start index

   Using enumerate(dataloader, 1) starts counting batches at 1, so modulo 5 works correctly for every 5 batches.

2. Step 2: Verify condition and reset

   The condition i % 5 == 0 triggers every 5 batches. After printing average loss, running_loss resets to 0.0 to accumulate next 5 batches.

3. Step 3: Compare other options

   for epoch in range(num_epochs): running_loss = 0.0 for i, data in enumerate(dataloader): inputs, labels = data optimizer.zero_grad() outputs = model(inputs) loss = loss_fn(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() if i % 5 == 0: print(f'Batch {i}, Avg Loss: {running_loss / 5:.4f}') running_loss = 0.0 starts enumerate at 0, so i % 5 == 0 triggers at batch 0 (nonexistent). for epoch in range(num_epochs): running_loss = 0.0 for i, data in enumerate(dataloader, 1): inputs, labels = data optimizer.zero_grad() outputs = model(inputs) loss = loss_fn(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() if i % 4 == 0: print(f'Batch {i}, Avg Loss: {running_loss / 5:.4f}') running_loss = 0.0 uses i % 4 == 0, wrong interval. for epoch in range(num_epochs): running_loss = 0.0 for i, data in enumerate(dataloader, 1): inputs, labels = data optimizer.zero_grad() outputs = model(inputs) loss = loss_fn(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() if i % 5 == 1: print(f'Batch {i}, Avg Loss: {running_loss / 5:.4f}') running_loss = 0.0 triggers at i % 5 == 1, which is off by one.

4. Final Answer:

   Option A correctly prints average loss every 5 batches with proper indexing and reset. -> Option A

5. Quick Check:

   enumerate start=1 and i % 5 == 0 for every 5 batches [OK]

Hint: Start enumerate at 1 and use i % 5 == 0 to print every 5 batches [OK]

Common Mistakes:

• Starting enumerate at 0 causing off-by-one errors
• Using wrong modulo condition like i % 4 == 0
• Resetting running_loss at wrong time