Writing a Training Loop in PyTorch

Moving a training loop to a GPU changes where the data and model live, not the training logic itself.

Programming

A Familiar Loop on New Hardware

You have already seen a complete training loop that runs on the CPU. It creates a model, feeds it batches of features and labels, computes a loss, and updates the parameters with an optimizer. This article answers a practical question: if a GPU is available, what actually has to change in that loop, and what should stay the same?

The short answer is that the training logic does not change. The loop still performs the same forward pass, loss computation, gradient reset, backward pass, and optimizer step. What changes is where the numbers live while those operations run. Keep that distinction in mind as you read on, because most GPU-related errors come from confusing the logic of the loop with the location of its data.

Which Lines Change and Which Stay Put

The GPU version of the loop is the CPU version with three edits. The first creates a device object that names the hardware to use. The second moves the model to that device once, before training starts. The third moves each batch to that device inside the inner loop. Everything else in the loop is untouched.

Which lines in a CPU training loop move to the GPU (model, batch, targets), and which lines stay unchanged because the training logic itself is the same?

Read the diagram as a checklist. The highlighted steps are the only ones that carry device information. The forward pass, the loss computation, the gradient reset, the backward pass, and the optimizer step appear unchanged in both versions. If your GPU version of the loop differs from the CPU version in any other way, something other than the three edits has changed.

When you run the GPU version, compare its printed output with the CPU run. The printed format includes the epoch, the batch index, and the loss. The output should look nearly identical to the CPU output, and the loss values should follow the same pattern. Moving computation to a different device does not change the mathematics of training, so a large difference in the losses is a signal to look for a bug rather than a sign that the GPU is working differently.

Where the Model and Batches Must Meet

Two streams of data must meet on the GPU before any computation can happen. The first is the model parameters, which move once at the start of training. The second is each batch of features and labels, which moves on every iteration because batches are produced fresh each time through the loop. A training step can run only when both have arrived on the same device.

Where does the model's parameters live and where does each batch of data live, and what breaks when they sit on different devices during the forward pass?

The two transfer lines appear in different places in the code, even though they serve the same purpose. The model transfer happens once, right after the model is created. The batch transfer happens inside the inner loop, right after each batch is loaded. If either transfer is missing, or if the two point at different devices, the forward pass has inputs and parameters that do not live together, and the step cannot run as intended.

Selecting the Device Once

The device object is the single place where you say which hardware to use. Once it exists, the rest of the code refers to it. Hard-coding the string cuda would make the script fail on a computer without an NVIDIA GPU. The recommended pattern selects cuda when a GPU is available and falls back to the CPU otherwise. This matters most when you share PyTorch code with others, because the same file then works on many machines.

What path does the code take when a CUDA GPU is available versus when it is not, and which device object is ultimately returned?

The following code shows the pattern applied to the loop from earlier. The device line is the only place that decides the hardware. The model is moved once, and each batch is moved inside the inner loop. The rest of the training logic, which is elided here with a comment, is the same as in the CPU version.

python

import torch
import torch.nn as nn

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

model = nn.Linear(4, 1).to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()

for epoch in range(3):
    for batch_index, (features, labels) in enumerate(loader):
        features = features.to(device)
        labels = labels.to(device)
        # forward pass, loss, gradient reset, backward pass, optimizer step
        # are the same lines as in the CPU version
        ...

import torch
import torch.nn as nn

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

model = nn.Linear(4, 1).to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()

for epoch in range(3):
    for batch_index, (features, labels) in enumerate(loader):
        features = features.to(device)
        labels = labels.to(device)
        # forward pass, loss, gradient reset, backward pass, optimizer step
        # are the same lines as in the CPU version
        ...

Notes on the code

The device object is created once and selects cuda when available, otherwise cpu
The model is moved once, before the loop starts
Each batch is moved inside the inner loop because batches are produced fresh on every iteration
The training logic between the transfers is unchanged

The source also notes that you can skip the device object and call .to('cuda') directly on a tensor. The result is the same, and the line is shorter. Either style works, but the device-object style makes it easier to reuse one variable across the model and every batch. Because .to(device) behaves the same whether device points to a GPU or the CPU, the rest of the training loop does not need to know which branch was taken.

Following One Training Step

Trace a single iteration of the inner loop in order. Each step below either happens on the host side, where batches are loaded and handed over, or on the accelerator side, where the arithmetic runs. The sequence is the same for every batch, which is why the transfer must happen on every pass.

1
Load the next batch of features and labels from the data loader on the host side
2
Move the features to the device with .to(device)
3
Move the labels to the device with .to(device)
4
Run the forward pass, where the model and the batch now share the same device
5
Compute the loss from the outputs and the labels, both of which live on the device
6
Reset the gradients, run the backward pass, and apply the optimizer step on the device

Worked example

Tracing a Device Mismatch Through One Step

Suppose the device object is created correctly and the model is moved to it, but the batch transfer line is accidentally left out of the inner loop. Where does the step break down?

1
Check the model
The model parameters were moved once at the start, so they live on the device.
2
Check the batch
The batch was loaded on the host side and never moved. It still lives on the CPU.
3
Check the forward pass
The forward pass needs the model and its input on the same device. Here they sit on different devices, so the step cannot run as intended.
4
Fix the cause
Add the .to(device) call for the batch inside the inner loop. The model transfer and the training logic stay exactly as they were.

Result

The bug is a missing batch transfer, not a problem with the training logic. The fix is one line inside the inner loop.

When the GPU Pays Off

A natural expectation is that moving training to a GPU always makes it faster. The source is more careful. For the small training loop used in the lesson, you probably will not see a speedup, because transferring the model and the data from the CPU to the GPU costs time. The GPU must receive the model and the data before it can work on them, and for a small model that transfer can take as long as the computation it saves.

The picture changes for larger workloads. The source says you can expect a significant speedup when training deep neural networks, especially large language models. In those cases the arithmetic is heavy enough that the one-time model transfer and the per-batch transfers become a small fraction of total training time.

Factor	Pushes toward a speedup	Pushes toward no speedup
Model size	Deep neural networks and large language models with heavy arithmetic	A small model where the computation saved is about the same as the transfer cost
Transfer cost	Transfer is a small fraction of total training time	Model and data copies take as long as the work they enable
Batch work	Each batch carries enough computation to amortize its transfer	Each batch carries little computation relative to moving it
Likely outcome	A significant speedup on the GPU	The GPU run may be no faster than the CPU run

Factors that decide whether the GPU run is likely to be faster

Mistakes That Cause Device Errors

Common mistake
Hard-coding the cuda device
What it looks like
```
device = torch.device('cuda') on a laptop with no NVIDIA GPU
```
Why it fails
The script fails on any computer without an NVIDIA GPU, so the code is not portable.
Do this instead
Select the device with torch.device('cuda' if torch.cuda.is_available() else 'cpu').
Common mistake
Moving the model but not the batches
What it looks like
model.to(device) is called, but features and labels are used directly inside the inner loop
Why it fails
The model's parameters and the batch then live on different devices when the forward pass runs.
Do this instead
Call .to(device) on the features and labels inside the inner loop, every iteration.
Common mistake
Moving the batch once before the loop
What it looks like
features and labels are moved before the for loop starts and reused
Why it fails
Batches are produced fresh on every iteration, so the data in the loop is new each time and must be moved each time.
Do this instead
Move each batch inside the inner loop, right after it is loaded.
Common mistake
Assuming the GPU is always faster
What it looks like
Expecting a tiny model on a GPU to finish much sooner than on a CPU
Why it fails
Transfer cost can cancel the benefit for small models.
Do this instead
Estimate whether the model and batches are large enough that computation dominates transfer before expecting a speedup.
Common mistake
Changing the training logic when moving to a GPU instead of only the data placement

Debugging

Suppose a GPU run prints losses that differ noticeably from a CPU run. Work through the loop in order. Check the device line first, because a wrong device propagates everywhere. Then check the model transfer. Then check the batch transfer inside the inner loop. Finally, check that the loss is computed from outputs and labels that both live on the same device.

Practice Moving a Loop to the GPU

MEDIUM

You have a working CPU training loop. Rewrite it so it runs on a GPU when one is available and on the CPU otherwise. List exactly which lines you changed, and confirm that the forward pass, loss, gradient reset, backward pass, and optimizer step are unchanged.

Hints

Start by creating a single device object that selects cuda when available and cpu otherwise
Move the model once, right after it is created
Move the features and the labels inside the inner loop, every iteration
If you changed any line outside those three places, look again at whether it was needed

Your attempt (not saved)

EASY

A colleague says a small model will always train faster on a GPU. Using the factors from this article, explain what you would check before agreeing, and what result you would expect for a very small model with tiny batches.