PyTorch Training Loops, DataLoaders, GPUs and Debugging Explained

Your PyTorch code runs without a single error. The loss prints a number every epoch. And yet the model learns nothing, or worse, it learns the wrong thing. In deep learning, the most dangerous bugs never crash. Today we build the habits that catch them.

• No error does not mean correct
• Silent bugs are the worst

In lesson 49 we met tensors and autograd, which compute gradients for us. Now we assemble the full workflow around them. We will define models, feed them data in batches, run a clean train and validate loop, move everything to a GPU, and finish with a debugging checklist you can use on any project.

• Module, Dataset, DataLoader
• Train loop, GPU, debugging

A model in PyTorch is a class that inherits from nn.Module. In the init method you create layers, and in the forward method you describe how an input flows through them. Any layer you assign as an attribute is registered automatically, so model dot parameters can hand every weight to the optimizer.

• __init__ builds layers
• forward defines the computation

When the data just flows straight through, nn.Sequential saves you writing a class. Take a network from four inputs to eight hidden units to three outputs, with a ReLU between. The first layer has four times eight weights plus eight biases, which is forty. The second has eight times three plus three, which is twenty seven. Total, sixty seven parameters.

• Sequential chains layers
• Count: weights plus biases

A loss function turns predictions and targets into one number to minimize. For classification, nn.CrossEntropyLoss expects raw scores called logits, and integer class labels. It applies log softmax internally. A very common mistake is adding your own softmax at the end of the model, which squashes the signal and hurts training. For regression, use nn.MSELoss.

• CrossEntropyLoss wants raw logits
• Do not apply softmax first

The optimizer owns the update rule. You create it with the model parameters and a learning rate, for example SGD or Adam. Every iteration follows the same beat: zero the gradients, call backward on the loss, then call step. Forgetting to zero the gradients is a classic bug, because PyTorch adds new gradients onto old ones.

• zero_grad, backward, step
• Gradients accumulate by default

A Dataset answers just two questions. How many samples do you have, through the len method, and what is sample number i, through the getitem method. That is the whole contract. Inside getitem you might read a file, apply a transform, and return a feature tensor with its label. Keep it focused on one sample at a time.

• __len__ returns the size
• __getitem__ returns one sample

The DataLoader wraps a dataset and serves it in batches. Batching gives the GPU enough parallel work, and averages the gradient over several samples, which makes updates less noisy. Suppose you have one thousand samples and a batch size of thirty two. That gives thirty one full batches and one last batch of eight, so thirty two steps per epoch.

• Groups samples into batches
• Why batches beat single samples

For training, set shuffle to true, so each epoch sees the samples in a new order and batches do not become repetitive. For validation, leave it false, since order does not matter and you want stable results. If that small final batch causes trouble, for example with batch statistics, drop_last true discards it, leaving thirty one steps.

• Shuffle training data only
• drop_last skips the small batch

Loading data can be slower than the model itself. Setting num_workers above zero starts separate processes that prepare upcoming batches while the GPU computes. Setting pin_memory true places batches in page locked memory, which makes copies to the GPU faster. More workers is not always better, so measure on your own machine.

• num_workers loads in parallel
• pin_memory speeds host to GPU copy

By default, the loader stacks samples into one tensor, which requires identical shapes. Sentences of different lengths break that rule. A custom collate function receives the list of samples and decides how to combine them. A typical choice is padding every sentence to the longest one in the batch, and returning a mask that marks the real tokens.

• Default collate stacks tensors
• Custom collate pads sequences

Here is the heart of every project. For each epoch, loop over batches. Move the batch to the device, run the forward pass, compute the loss, zero the gradients, call backward, and step the optimizer. Those few lines are the same whether the model has sixty seven parameters or sixty seven billion.

• For each epoch, for each batch
• Forward, loss, backward, step

Some layers behave differently in training and in evaluation. Dropout randomly zeroes units while training, and batch normalization uses batch statistics. Calling model dot eval switches them to their deterministic inference behavior, and model dot train switches back. Forgetting eval during validation gives noisy, pessimistic numbers, as lesson 48 hinted.

• Dropout and batch norm differ
• Eval turns randomness off

Validation must not learn. Wrap it in torch dot no_grad, so PyTorch skips building the graph, saving memory and time. To report the loss, multiply each batch loss by its batch size, add it up, and divide by the total samples. With a last batch of eight, this weighting keeps that small batch from counting as much as a full one.

• torch.no_grad saves memory
• Average loss weighted by batch size

When logging, store loss dot item, a plain Python number. If you append the loss tensor itself to a list, each one keeps its whole computation graph alive, and memory quietly grows until you run out. Log training loss, validation loss, and a metric such as accuracy every epoch. The gap between training and validation tells you about overfitting.

• loss.item() gives a plain float
• Never keep the loss tensor

Read the curves like a doctor reads a chart. If both losses stay high, the model is underfitting, so give it more capacity or train longer. If training keeps falling while validation turns upward, it is overfitting, so use regularization, more data, or stop earlier. If the training loss is erratic, try a smaller learning rate.

• Both high: underfitting
• Gap widening: overfitting

A state dict is a dictionary mapping each layer name to its tensors. Saving it with torch dot save is the recommended way to store a model. For resuming training, also save the optimizer state, since Adam keeps running averages, plus the epoch number. To load, build the model first, then call load_state_dict.

• Save the state_dict, not the object
• Include optimizer and epoch

Two details trip people up when loading. First, a checkpoint saved on a GPU needs map_location set, if you load it on a machine without one. Second, after loading for inference, call eval, or dropout will still be active. Also, treat checkpoint files from strangers carefully, because loading can execute code unless you use the safer weights only option.

• Call model.eval() after loading
• Use map_location for devices

Early stopping watches the validation loss and halts when it stops improving. Say the best validation loss was zero point forty two at epoch seven, and your patience is three. If epochs eight, nine and ten fail to beat it, training stops, and you restore the epoch seven checkpoint. You keep the best model, not the last one.

• Track best validation loss
• Stop after patience epochs

Neural networks are mostly matrix multiplications, and each output element can be computed independently. A GPU has thousands of simple cores built for exactly that kind of parallel arithmetic. A CPU has a few powerful cores that excel at varied, sequential work. For large models the difference can be enormous, though tiny models may not benefit at all.

• Thousands of parallel cores
• Matrix multiplies dominate

The standard pattern is to pick a device once, using cuda if available, otherwise cpu. Then move the model with model dot to device, and move every batch of inputs and targets too. The classic error message says tensors are on different devices. The fix is nearly always a forgotten dot to on the model, the inputs, or the targets.

• device = cuda if available
• Model and data on the same device

Remember pin_memory from the DataLoader. When batches sit in pinned memory, you can copy them with non_blocking set to true. The copy then happens asynchronously, so the GPU can still be finishing the previous step. It is a small change, but on input heavy workloads it can reduce idle time.

• Pinned pages copy faster
• non_blocking overlaps copy and compute

Mixed precision runs most operations in sixteen bit numbers, which use half the memory and run faster on modern GPUs. Wrapping the forward pass in torch autocast lets PyTorch choose the precision per operation, keeping numerically sensitive ones in thirty two bit. The loss and parameter updates remain in full precision, so accuracy is usually preserved.

• Half precision is faster, smaller
• Keep sensitive ops in float32

Float sixteen cannot represent very small numbers, so tiny gradients can underflow to zero. A GradScaler multiplies the loss by a large factor before backward, so gradients stay representable, then divides them back before the optimizer step. If it sees infinities, it skips that step and lowers the scale. The bfloat16 format has a wider range and usually needs no scaler.

• float16 tiny values round to zero
• Scale loss up, unscale gradients

When you hit an out of memory error, reduce the batch size first, since stored activations grow with it. If you still want a larger effective batch, use gradient accumulation. Run four mini batches of sixteen, divide each loss by four, and step the optimizer once. The result approximates a batch of sixty four. Also delete unused tensors, and avoid holding on to graphs.

• Activations dominate memory
• Gradient accumulation fakes big batches

Training uses randomness for weight initialization, shuffling, dropout, and augmentation. To repeat an experiment, seed every source: Python's random module, NumPy, and torch dot manual_seed. With the same seed, the same initial weights appear every run. This matters because a difference between two models might just be luck from different seeds.

• Seed Python, NumPy and torch
• Randomness enters in many places

Seeds alone do not guarantee identical results. Some GPU operations are nondeterministic for speed, and you can request deterministic behavior with torch dot use_deterministic_algorithms, at some cost in speed. With DataLoader workers, pass a seeded generator and a worker init function, so every worker shuffles and augments reproducibly. Perfect bitwise reproduction across different hardware is still not promised.

• use_deterministic_algorithms
• Seed each worker via a generator

Start debugging by printing the shape of one batch and one output. A very sneaky bug appears when predictions have shape thirty two by one and targets have shape thirty two. Subtracting them broadcasts into a thirty two by thirty two matrix, and the loss runs without error but is wrong. Always confirm that shapes match what you intend.

• Print shapes at each stage
• Watch silent broadcasting

Before training, check the loss at initialization. With ten classes and a model that guesses roughly uniformly, cross entropy should be about the natural log of ten, which is two point three. If you see a loss of forty, or zero point zero, something is off, perhaps a wrong label range, an extra softmax, or badly scaled inputs. A wrong first number is a free clue.

• Expect log of number of classes
• Ten classes: about 2.30

This is the single most useful test. Take one small batch, say eight samples, and train on it repeatedly. A healthy model with enough capacity should drive the loss close to zero within a few hundred steps. If it cannot memorize eight examples, the bug is in the model, the loss, the labels, or the optimizer, not in your data volume.

• Train on 1 batch repeatedly
• Loss must head toward zero

Watching gradient norms reveals problems that the loss hides. After backward, compute the overall norm of all gradients. A norm near zero hints at vanishing gradients, as in lesson 47, and a huge norm hints at explosion. The function clip_grad_norm_ rescales gradients whose norm exceeds a limit, such as one, so a single bad batch cannot wreck the weights.

• Compute total gradient norm
• clip_grad_norm_ tames explosions

When the loss becomes NaN, find where it first appears. Turn on torch autograd set_detect_anomaly to get a traceback pointing at the operation, and check inputs with torch isnan. Typical causes are a learning rate that is too high, taking the log of zero, dividing by zero, or corrupt data. Once one NaN appears, it spreads everywhere.

• Find the first bad value
• Common causes: lr, log, division

Put it together as a routine. First verify data and shapes. Second, check the initial loss. Third, overfit one batch. Only then train on the full dataset, while logging losses and gradient norms. Change one thing at a time. This order moves from the cheapest checks to the most expensive, and it saves days.

• Shapes, init loss, one batch
• Then scale up and tune

Let us recap. Modules hold parameters, losses and optimizers drive learning, and Dataset plus DataLoader deliver batches. The loop has a fixed beat, with train and eval modes, no_grad validation, checkpoints and early stopping. Move model and data to the device, use autocast carefully, seed everything, and debug by checking shapes, initial loss, and overfitting one batch.

• Module, Dataset, DataLoader
• Loop, checkpoint, GPU, debug

You now have a complete workflow, and in the next lesson we put it to work on images. Lesson 51 explains convolution mechanics: kernels, stride, padding and pooling.

• Lesson 51: convolution mechanics
• Subscribe for the full course