How to run a training loop in PyTorch

A PyTorch training loop is the code that turns mini-batches into model parameter updates. Keeping the loop explicit helps when a project needs custom data handling, loss calculations, logging, or optimizer timing instead of a higher-level training framework.

A minimal loop needs a model, a DataLoader, a loss function, and an optimizer. For each batch, the loop resets stale gradients, runs the forward pass, computes loss, calls backward(), and lets the optimizer update the parameters.

The smoke script uses a tiny regression dataset on CPU so the training signal is easy to inspect. Lower epoch loss, weight_changed=True, and a prediction tensor shape confirm that the loop executed the training pass and the trained model still returns outputs.

1. Create a training-loop smoke script.

   training_loop_demo.py

   import torch
   from torch import nn
   from torch.utils.data import DataLoader, TensorDataset
    
    
   torch.manual_seed(23)
    
   features = torch.tensor(
       [
           [-1.0, 0.0, 0.5],
           [-0.5, 0.25, 1.0],
           [0.0, -0.5, 0.25],
           [0.5, 0.75, -0.25],
           [1.0, -0.25, -0.5],
           [1.5, 0.5, 0.0],
           [2.0, -0.75, 0.75],
           [2.5, 1.0, -1.0],
       ],
       dtype=torch.float32,
   )
   targets = features @ torch.tensor([[0.8], [-0.4], [0.3]]) + 0.2
    
   train_data = TensorDataset(features, targets)
   train_loader = DataLoader(train_data, batch_size=4, shuffle=False)
    
   model = nn.Sequential(nn.Linear(3, 1))
   loss_fn = nn.MSELoss()
   optimizer = torch.optim.SGD(model.parameters(), lr=0.08)
    
   first_weight = model[0].weight
   weight_before = first_weight.detach().clone()
    
   model.train()
   epoch_losses = []
    
   for epoch in range(1, 7):
       loss_total = 0.0
       sample_count = 0
    
       for batch_features, batch_targets in train_loader:
           optimizer.zero_grad(set_to_none=True)
           predictions = model(batch_features)
           loss = loss_fn(predictions, batch_targets)
           loss.backward()
           optimizer.step()
    
           loss_total += loss.item() * batch_features.size(0)
           sample_count += batch_features.size(0)
    
       average_loss = loss_total / sample_count
       epoch_losses.append(average_loss)
       print(f"epoch={epoch} loss={average_loss:.6f}")
    
   with torch.inference_mode():
       sample_prediction = model(features[:2])
    
   weight_changed = not torch.equal(weight_before, first_weight.detach())
   loss_decreased = epoch_losses[-1] < epoch_losses[0]
    
   print(f"weight_changed={weight_changed}")
   print(f"loss_decreased={loss_decreased}")
   print(f"prediction_shape={tuple(sample_prediction.shape)}")

   shuffle=False keeps the smoke output repeatable. Use shuffle=True for ordinary training data when batch order should vary between epochs.

2. Run the training-loop smoke script.

   $ python training_loop_demo.py
   epoch=1 loss=1.753747
   epoch=2 loss=0.223522
   epoch=3 loss=0.080783
   epoch=4 loss=0.053622
   epoch=5 loss=0.040088
   epoch=6 loss=0.030757
   weight_changed=True
   loss_decreased=True
   prediction_shape=(2, 1)

3. Confirm the training proof lines.

   loss_decreased=True means the final epoch loss is lower than the first epoch loss. weight_changed=True confirms that optimizer.step() updated the linear layer, and prediction_shape=(2, 1) confirms that the trained model still returns two one-value predictions.

4. Move the same batch-loop order into the project training function.

   def train_one_epoch(model, train_loader, loss_fn, optimizer, device):
       model.train()
       loss_total = 0.0
       sample_count = 0
    
       for batch_features, batch_targets in train_loader:
           batch_features = batch_features.to(device)
           batch_targets = batch_targets.to(device)
    
           optimizer.zero_grad(set_to_none=True)
           predictions = model(batch_features)
           loss = loss_fn(predictions, batch_targets)
           loss.backward()
           optimizer.step()
    
           loss_total += loss.item() * batch_features.size(0)
           sample_count += batch_features.size(0)
    
       return loss_total / sample_count

   optimizer.zero_grad(set_to_none=True) clears stale gradients before each backward pass. Move every batch tensor to the same device as the model before the forward pass.
   Related: How to zero gradients in PyTorch
   Related: How to select a device in PyTorch

5. Call the training function for each epoch and log the average loss.

   for epoch in range(1, epochs + 1):
       train_loss = train_one_epoch(
           model,
           train_loader,
           loss_fn,
           optimizer,
           device,
       )
       print(f"epoch={epoch} train_loss={train_loss:.6f}")

   Add evaluation, checkpointing, or scheduler calls outside the batch loop unless the project needs per-batch behavior.
   Related: How to run evaluation in PyTorch
   Related: How to save and restore a PyTorch training checkpoint
   Related: How to set a learning rate scheduler in PyTorch

6. Remove the temporary smoke script.

   $ rm training_loop_demo.py