Measuring Training Performance on Different Hardware

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

Programming

Same Loop, Different 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 asks a practical question: if a GPU is available, what actually has to change in that loop, and what should stay the same?

What do you think happens?

If you move a training loop from the CPU to a GPU, does the training logic itself change?

Yes, the forward pass and the loss are rewritten for the GPU
No, the training logic stays the same and only the location of the data and model changes

Reveal answer

Answer:

No, the training logic stays the same and only the location of the data and model changes.

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 this distinction in mind throughout the article. Most GPU-related errors come from confusing the logic of the loop with the location of its data. A line of code can be correct as logic and still fail because the tensor it touches sits on the wrong device.

What Stays the Same in the Loop

Every operation in a CPU training loop has an exact counterpart in a GPU training loop. The GPU version performs the same steps in the same order: a forward pass that turns a batch into predictions, a loss computed from those predictions and the labels, a gradient reset, a backward pass that computes gradients, and an optimizer step that updates the parameters.

Because the mathematics does not change, the printed results should look nearly the same. The source reports that GPU results are similar to those obtained on the CPU, which is what you should expect: moving computation to a different device does not change the mathematics of training.

Part of the loop	Changes on a GPU?	What it means in practice
Device selection	Yes, new	A device object is created that names the hardware to use
Model	Yes, modified	The model is moved to the device once, before training starts
Each batch of features and labels	Yes, modified	Each batch is moved to the device inside the inner loop, every iteration
Forward pass	No	Same operations, now running where the data lives
Loss computation	No	Computed from logits and labels that must share the device
Gradient reset	No	Same call, same purpose
Backward pass	No	Same call, gradients computed on the device holding the data
Optimizer step	No	Same update of the parameters

The three edits to a CPU loop are device selection, model transfer, and batch transfer. Every other line keeps its logic.

Tracing Where the Data Lives

The GPU version of the loop is the CPU version with three edits. As you read the code, trace three things: where the device object is created, where the model gets moved, and where each batch gets moved inside the inner loop. 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.

The source creates the device object with the string cuda, which names the NVIDIA GPU, and then moves the model to it. The model is moved once, because its parameters persist from one iteration to the next. Batches are different: they are produced fresh on every iteration, so each one must be moved every time it is used.

Which data must be on the GPU before a training step can run, and when does each one arrive? The model moves once at the start, while each batch moves on every iteration.

The diagram shows the two streams of data that must meet on the GPU. A training step can only run when both have arrived on the same device. That is why the transfer lines appear in different places in the code, even though they serve the same purpose: one moves the model before the loop begins, and the other moves each batch inside it.

Inside One Iteration

The inner loop is where most transfer activity happens, so it helps to read one iteration in order. The sequence below lists what happens on each pass, with the batch move placed before any computation that uses the batch.

What happens in order on each iteration? The batch is copied to the device first, then the forward pass, loss, gradient reset, backward pass, and optimizer step run on that device.

Only the move step is new. Every other step is the same code you already wrote for the CPU. When the next iteration begins, the loader produces a fresh batch on the CPU, and the move step runs again for that batch.

Selecting the Device With a Fallback

Hard-coding cuda makes the script fail on a computer without an NVIDIA GPU. The source recommends a more portable pattern: select the GPU when one is available and fall back to the CPU otherwise. This is the recommended approach when sharing PyTorch code with others, because the same file then works on many machines.

Which branch does the program take with and without a GPU, and how does the chosen device reach the model and the data?

Code walkthrough

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)

        logits = model(features)
        loss = loss_fn(logits, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print(epoch, batch_index, loss.item())

1
The import lines are the same as in the CPU version.
2
The device line is the single place where the hardware is chosen. It selects cuda when torch.cuda.is_available() reports a GPU and falls back to cpu otherwise.
3
The model is moved once with .to(device) before training starts. Its parameters then live on the chosen device.
4
The optimizer and loss function are created exactly as in the CPU version.
5
Inside the inner loop, each batch is moved to the device. This happens every iteration because each batch is produced fresh.
6
The forward pass, loss, gradient reset, backward pass, and optimizer step are unchanged. The loss uses logits and labels that both live on the same device.
7
The print statement shows the epoch, batch index, and loss, so you can compare its output with the CPU run.

This example is illustrative: it assumes the loader, the loss function, and the optimizer are defined as they were in the CPU version. 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, but the device-object style makes it easier to reuse one variable across the model and every batch.

Mistakes That Split Model and Data

Most GPU training failures come from one of three places. The model and data end up on different devices, the device choice is wrong, or the transfer is applied in the wrong place. The items below describe each mistake and how to correct it.

Common mistake
Moving the model to the GPU but leaving the batches on the CPU
What it looks like
```
model = nn.Linear(4, 1).to(device) is written, but the loop uses features directly without calling .to(device) on it.
```
Why it fails
A training step needs the model and each batch on the same device. When they are split, the forward pass cannot combine them.
Do this instead
Move each batch inside the inner loop with features = features.to(device) and labels = labels.to(device).
Common mistake
Hard-coding cuda in the device line
What it looks like
```
device = torch.device('cuda') on a computer without an NVIDIA GPU.
```
Why it fails
The script fails on any machine that lacks an NVIDIA GPU, which makes the code hard to share.
Do this instead
Use torch.device('cuda' if torch.cuda.is_available() else 'cpu') so the same file runs on both kinds of machine.
Common mistake
Moving the batch once before the loop
What it looks like
features and labels are moved to the device above the for loop, and the loop reuses them.
Why it fails
Batches are produced fresh on every iteration, so a batch moved before the loop does not cover the batches the loader produces later.
Do this instead
Move each batch inside the inner loop, immediately after it is produced.
Common mistake
Assuming a loss difference means the logic changed
What it looks like
The GPU run prints losses that differ noticeably from the CPU run, and the loop logic is rewritten.
Why it fails
The source expects GPU results to be similar to CPU results. A large difference usually points to a device or transfer problem, not to the training logic.
Do this instead
Check the device line first, then the model transfer, then the batch transfer inside the inner loop, and finally confirm that logits and labels share the device.

When the Speedup Appears

A natural expectation is that moving training to a GPU always makes it faster. The source is more careful. For the small training loop in its example, it says you will probably not see a speedup, because transferring the model and 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.

Workload	Expected GPU effect	Why
Small model, small batches	Little or no speedup, possibly slower	Transfer cost of the model and data can match or exceed the computation saved
Deep neural networks	Significant speedup	Heavy arithmetic makes the one-time and per-batch transfer costs a small fraction of total training time
Large language models	Significant speedup	The same reasoning applies at a larger scale, where arithmetic dominates

Whether a GPU helps depends on how the cost of computation compares with the cost of moving data.

Predicting Where the Batch Goes

MEDIUM

A colleague moves the model to the GPU before the training loop starts, and inside the inner loop they compute the forward pass on features that were loaded on the CPU. Predict what happens when the first training step runs on a machine with a GPU, and explain which line you would change.

Hints

Ask which data must be on the same device as the model before the forward pass can run.
Think about whether each batch is produced once or fresh on every iteration.
The fix belongs inside the inner loop, not before it.

Your attempt (not saved)

EASY

Write the device line and the two transfer lines for a model and a loader, so that the same script runs on a GPU when present and on the CPU otherwise.