You trained your first network on a toy set that fit in one tensor, and it felt easy. Then real data arrives, a CSV with thousands of rows or an image folder too big for memory, and the PyTorch DataLoader is the tool built for exactly that. Here you pair it with a custom Dataset, feed a network from a CSV and an image folder, run the five-step loop from memory, validate honestly, and save a checkpoint you can resume from.
“The most common neural net mistakes: you forgot to toggle train/eval mode for the net, and you forgot to zero_grad() before backward().”
Andrej Karpathy, on training neural networks
Last Updated: July 2026 | Tested on: Python 3.14.6, PyTorch 2.12.1 | Difficulty: Intermediate | Reading Time: 17 minutes
In the first PyTorch tutorial we trained on the tiny Iris set with a one-line TensorDataset. That is fine for a demo, but your real data lives in CSV files, image folders, and databases, and it rarely fits in memory all at once. This is the job the Dataset and DataLoader pair was built for. The Dataset knows how to fetch one sample, and the DataLoader turns those single samples into shuffled batches, so your training loop stays short no matter where the data comes from.
Here is what we cover:
- Writing a custom Dataset with
__len__and__getitem__ - DataLoader batching and shuffling, plus image transforms
- The five-step training loop and why
zero_grad()matters - Validation with
model.train(),model.eval(), andtorch.no_grad() - Checkpointing: save, load, and resume mid-run
Table of Contents
Prerequisites
- First neural network with PyTorch (tensors, autograd, nn.Module)
- pip install torch torchvision
- A folder to work in; every script below runs on a plain CPU
Your Own Dataset: __len__ and __getitem__
Think of a Dataset like a well-organised pantry. You do not tip every jar onto the counter at once. You ask for one jar by its shelf number and the pantry hands it over. A PyTorch Dataset works the same way: it answers two questions only. How many samples are there (__len__), and give me sample number i (__getitem__). PyTorch handles the rest.
Say a grocery startup run by a developer named Aditi wants to sort produce into fresh, aging, and spoiled from four measurements: weight, firmness, sugar percent, and days since harvest. The labels live in a plain produce.csv with 240 rows. Like the other project posts in this series, we generate that data in code so every block here runs on any machine, no download needed. This one script writes both files the post uses: produce.csv and the veg_photos folder the ImageFolder demo reads later. Run it once.
📄 make_data.py: generate produce.csv and the veg_photos folder
# make_data.py: synthesize produce.csv and the veg_photos folder so every
# script in this post runs on any machine, no download needed.
import csv, os
import numpy as np
from PIL import Image
rng = np.random.default_rng(42)
# --- produce.csv: 240 rows, 80 per freshness class ---
# label 0 = fresh, 1 = aging, 2 = spoiled
specs = [ # (label, weight mean, firmness mean, sugar mean, days mean)
(0, 150, 8.5, 9.0, 2.0),
(1, 135, 5.5, 7.0, 7.0),
(2, 120, 2.5, 5.5, 12.0),
]
rows = []
for label, w, f, s, d in specs:
for _ in range(80):
rows.append([round(float(rng.normal(w, 12)), 1),
round(float(np.clip(rng.normal(f, 0.8), 0.5, 10)), 2),
round(float(np.clip(rng.normal(s, 0.9), 0.5, 15)), 2),
round(float(np.clip(rng.normal(d, 1.2), 0.5, 20)), 1),
label])
rng.shuffle(rows)
with open("produce.csv", "w", newline="") as fh:
wr = csv.writer(fh)
wr.writerow(["weight_g", "firmness", "sugar_pct", "days_since_harvest", "label"])
wr.writerows(rows)
print(f"Wrote produce.csv: {len(rows)} rows")
# --- veg_photos: 12 small colour images, one subfolder per class ---
colours = {"carrot": (237, 145, 33), "spinach": (60, 130, 53), "tomato": (220, 48, 35)}
for name, (r, g, b) in colours.items():
os.makedirs(f"veg_photos/{name}", exist_ok=True)
for i in range(4):
base = np.stack([np.full((64, 64), c, dtype=np.float64) for c in (r, g, b)], axis=-1)
noisy = np.clip(base + rng.normal(0, 25, base.shape), 0, 255).astype(np.uint8)
Image.fromarray(noisy).save(f"veg_photos/{name}/{name}_{i}.png")
print("Wrote veg_photos/: 3 classes x 4 images")
▶ Output
Wrote produce.csv: 240 rows Wrote veg_photos/: 3 classes x 4 images
The three freshness classes sit in well-separated clusters (fresh produce is heavier and firmer, spoiled produce is soft and old), which keeps the training demos later in the post fast and stable. The photo folder is just twelve flat-colour squares with noise, one subfolder per vegetable, which is all ImageFolder needs to show its directory convention. Here is the Dataset that reads the CSV.
📄 dataset_csv.py: a custom Dataset over a CSV file
import csv, torch
from torch.utils.data import Dataset, DataLoader
class ProduceDataset(Dataset):
"""Reads a CSV of produce features and freshness labels."""
def __init__(self, csv_path):
self.rows = []
with open(csv_path) as f:
reader = csv.DictReader(f)
for r in reader:
feats = [float(r["weight_g"]), float(r["firmness"]),
float(r["sugar_pct"]), float(r["days_since_harvest"])]
self.rows.append((feats, int(r["label"])))
def __len__(self):
return len(self.rows)
def __getitem__(self, idx):
feats, label = self.rows[idx]
x = torch.tensor(feats, dtype=torch.float32)
y = torch.tensor(label, dtype=torch.long)
return x, y
ds = ProduceDataset("produce.csv")
print(f"Dataset size: {len(ds)} samples")
x0, y0 = ds[0]
print(f"First sample features: {x0.tolist()} label: {y0.item()}")
loader = DataLoader(ds, batch_size=32, shuffle=True)
print(f"Batches per epoch: {len(loader)}")
xb, yb = next(iter(loader))
print(f"One batch -> X shape: {tuple(xb.shape)} y shape: {tuple(yb.shape)}")
print(f"Labels in this batch: {yb[:8].tolist()} ...")
▶ Output
Dataset size: 240 samples First sample features: [135.39999389648438, 3.380000114440918, 5.380000114440918, 10.5] label: 2 Batches per epoch: 8 One batch -> X shape: (32, 4) y shape: (32,) Labels in this batch: [1, 0, 1, 0, 0, 2, 2, 2] ...
What happened here: The class holds the raw rows and returns one (features, label) pair per index. That is the whole contract. Notice the label uses dtype=torch.long, because PyTorch’s classification loss expects integer class indices, not floats. We never wrote a batching loop, yet the DataLoader gave us 8 batches of 32 (240 divided by 32 is 7 full batches plus a last batch of 16), each already stacked into a neat (32, 4) tensor. The same Dataset would work if the rows came from a database or a JSON file. Only __getitem__ would change.
DataLoader: Batching, Shuffling, and Transforms
The PyTorch DataLoader is the conveyor belt that carries samples from the pantry to the model. You set three things that matter most: batch_size (how many samples per step), shuffle (reorder every epoch so the model does not memorise the row order), and transforms (per-sample preprocessing). For images, transforms do the heavy lifting: resize every photo to the same size, convert it to a tensor, and normalise the pixel values so training behaves.
You rarely write an image Dataset by hand. torchvision ships ImageFolder, which reads a directory where each subfolder is one class. Here it is on the small veg_photos folder that make_data.py drew earlier, one subfolder each for spinach, tomato, and carrot.
📄 dataset_images.py: ImageFolder with a transform pipeline
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# transforms: resize, turn into a tensor, normalise pixel values
transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(), # HxWxC uint8 -> CxHxW float in [0,1]
transforms.Normalize(mean=[0.5]*3, std=[0.5]*3),
])
ds = datasets.ImageFolder("veg_photos", transform=transform)
print(f"Classes found: {ds.classes}")
print(f"Class -> index: {ds.class_to_idx}")
print(f"Total images: {len(ds)}")
loader = DataLoader(ds, batch_size=4, shuffle=True)
images, labels = next(iter(loader))
print(f"Batch image tensor shape: {tuple(images.shape)}") # (B, C, H, W)
print(f"Pixel range after normalize: {images.min():.2f} to {images.max():.2f}")
print(f"Labels: {labels.tolist()}")
▶ Output
Classes found: ['carrot', 'spinach', 'tomato']
Class -> index: {'carrot': 0, 'spinach': 1, 'tomato': 2}
Total images: 12
Batch image tensor shape: (4, 3, 32, 32)
Pixel range after normalize: -0.91 to 0.88
Labels: [2, 2, 1, 2]
What happened here: ImageFolder scanned the directory, turned each subfolder name into a class, and mapped them to integer labels in alphabetical order (carrot 0, spinach 1, tomato 2). The transform pipeline ran on every image as it was loaded: ToTensor() reshaped it from height-width-channel bytes into the channel-first float layout PyTorch expects, and Normalize shifted the pixels from the 0 to 1 range into roughly -1 to 1, which is why the printed range is -0.91 to 0.88. The batch shape (4, 3, 32, 32) reads as four images, three colour channels, 32 by 32 pixels. That four-number shape is the standard image batch layout you will see everywhere.
Anatomy of One Epoch
Before we write the loop, hold the whole picture in your head. One epoch means one full pass over the data. The Dataset serves single samples, the PyTorch DataLoader groups them into shuffled batches, and for each batch you run the same five steps. After the batches are done, you switch to evaluation mode and check the model on held-out data, then the next epoch reshuffles and repeats.
Tip: click the diagram to open it full screen, then use the zoom and pan controls for a closer look.
The five inner steps never change, whether you are training a produce sorter or a language model: zero the old gradients, run the forward pass, measure the loss, run backward to fill in new gradients, and step the optimizer to update the weights. Memorise that order. The rest of this post is really just those five lines with guardrails around them.
Why zero_grad Matters
Here is the trap that bites almost everyone once. PyTorch does not overwrite gradients on each backward pass. It adds to them. Picture a kitchen scale you forget to reset between ingredients: every scoop reads heavier than it should because the last one is still counted. If you never call zero_grad(), your gradients pile up batch after batch and the optimizer takes wild, oversized steps. Let us watch it happen, then fix it.
📄 zero_grad_trap.py: gradients pile up without a reset
import torch
w = torch.tensor([2.0], requires_grad=True)
x = torch.tensor([3.0])
print("Calling backward 3 times WITHOUT zero_grad:")
for step in range(1, 4):
loss = (w * x - 5) ** 2 # gradient wrt w is 2*(wx-5)*x = 6 each time
loss.backward()
print(f" step {step}: w.grad = {w.grad.item():.1f}") # piles up: 6, 12, 18
print("\nNow zeroing the gradient before each backward:")
for step in range(1, 4):
if w.grad is not None:
w.grad.zero_()
loss = (w * x - 5) ** 2
loss.backward()
print(f" step {step}: w.grad = {w.grad.item():.1f}") # clean 6 every time
▶ Output
Calling backward 3 times WITHOUT zero_grad: step 1: w.grad = 6.0 step 2: w.grad = 12.0 step 3: w.grad = 18.0 Now zeroing the gradient before each backward: step 1: w.grad = 6.0 step 2: w.grad = 6.0 step 3: w.grad = 6.0
What happened here: The true gradient of this loss with respect to w is 6 every time. Without a reset it reads 6, then 12, then 18, because each backward call stacks onto the leftover. With zero_grad() (here shown as w.grad.zero_() so you see exactly what the optimizer does for you) it stays a clean 6. In a real loop you call optimizer.zero_grad() at the top of every batch and forget about it. The accumulation is not a bug, by the way. It is occasionally useful for simulating a bigger batch by summing gradients over several small batches before stepping. But that is opt-in, not the default you want.
Validation Done Right
Two switches keep validation honest, and forgetting either one produces confusing results. model.train() and model.eval() flip the mode of layers that behave differently during training, like dropout and batch normalization. torch.no_grad() tells autograd to stop building the graph, which saves memory and time when you are only measuring, not learning.
Our model uses dropout, which randomly zeros neurons while training to prevent overfitting. If you forget model.eval(), dropout stays on during inference and the same input gives a different answer every time. Watch it fail, then work. One heads-up before you run it: this script imports model_defs.py (defined in the Checkpointing section below) and loads checkpoint.pth (written by the training script in the last section), so read on and run train.py first if you want to execute it.
📄 eval_trap.py: dropout left on makes predictions wobble
import torch
from model_defs import ProduceDataset, FreshnessNet
torch.manual_seed(42)
ds = ProduceDataset("produce.csv")
model = FreshnessNet()
ckpt = torch.load("checkpoint.pth", weights_only=True)
model.load_state_dict(ckpt["model_state"])
sample, _ = ds[0]
sample = sample.unsqueeze(0) # one produce item, shape (1, 4)
# WRONG: still in train() mode, dropout randomly zeros neurons
model.train()
with torch.no_grad():
a = model(sample)
b = model(sample)
print("model.train() (dropout ON) - same input, two runs:")
print(f" run 1 logits: {a.squeeze().tolist()}")
print(f" run 2 logits: {b.squeeze().tolist()}")
print(f" identical? {torch.allclose(a, b)}")
# RIGHT: eval() turns dropout off, predictions are stable
model.eval()
with torch.no_grad():
a = model(sample)
b = model(sample)
print("\nmodel.eval() (dropout OFF) - same input, two runs:")
print(f" run 1 logits: {a.squeeze().tolist()}")
print(f" run 2 logits: {b.squeeze().tolist()}")
print(f" identical? {torch.allclose(a, b)}")
▶ Output
model.train() (dropout ON) - same input, two runs: run 1 logits: [-9.429152488708496, 0.8919063806533813, 2.990316390991211] run 2 logits: [-8.878755569458008, 0.6839334964752197, 3.9805140495300293] identical? False model.eval() (dropout OFF) - same input, two runs: run 1 logits: [-8.873634338378906, 0.1915067434310913, 4.134331703186035] run 2 logits: [-8.873634338378906, 0.1915067434310913, 4.134331703186035] identical? True
What happened here: In train() mode the same produce item gave two different logit vectors, because dropout knocked out a different random set of neurons each pass. That is helpful during training but makes evaluation meaningless. After model.eval() dropout switches off and the two runs are byte-for-byte identical. The rule: call model.train() before training batches and model.eval() before you validate or predict. And remember eval() and no_grad() are separate jobs, one controls layer behaviour and the other controls gradient tracking, so a proper validation pass uses both.
Checkpointing: Save, Load, and Resume
Training a big model can run for hours or days, and machines crash. A checkpoint is your save point in a long game. The key idea: to truly resume, you save more than the weights. You also save the optimizer state (Adam keeps running averages that matter) and the epoch number, so you pick up exactly where you left off instead of restarting the momentum from zero. We keep the model and dataset classes in a small model_defs.py so both the training script and the resume script import the same definitions.
📄 model_defs.py: shared Dataset and model definitions
import csv, torch
import torch.nn as nn
from torch.utils.data import Dataset
class ProduceDataset(Dataset):
def __init__(self, csv_path):
self.rows = []
with open(csv_path) as f:
for r in csv.DictReader(f):
feats = [float(r["weight_g"]), float(r["firmness"]),
float(r["sugar_pct"]), float(r["days_since_harvest"])]
self.rows.append((feats, int(r["label"])))
# standardize columns so the network trains smoothly
import statistics as st
cols = list(zip(*[f for f, _ in self.rows]))
self.mean = [st.mean(c) for c in cols]
self.std = [st.pstdev(c) for c in cols]
def __len__(self):
return len(self.rows)
def __getitem__(self, i):
feats, label = self.rows[i]
z = [(v - m) / s for v, m, s in zip(feats, self.mean, self.std)]
return torch.tensor(z, dtype=torch.float32), torch.tensor(label, dtype=torch.long)
class FreshnessNet(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(4, 24), nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(24, 3),
)
def forward(self, x):
return self.net(x)
This version standardises the four columns inside __getitem__ so the model trains smoothly, and it adds a dropout layer so the eval demo above had something to switch off. Now the resume script loads the checkpoint and continues.
📄 resume.py: reload weights and optimizer, then keep going
import torch, torch.nn as nn
from torch.utils.data import DataLoader
from model_defs import ProduceDataset, FreshnessNet # just the class defs
torch.manual_seed(42)
loader = DataLoader(ProduceDataset("produce.csv"), batch_size=32, shuffle=True)
model = FreshnessNet()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()
ckpt = torch.load("checkpoint.pth", weights_only=True)
model.load_state_dict(ckpt["model_state"])
optimizer.load_state_dict(ckpt["optim_state"])
start = ckpt["epoch"] + 1
print(f"Resumed from epoch {ckpt['epoch']} (last train loss {ckpt['train_curve'][-1]:.4f})")
model.train()
for epoch in range(start, start + 5):
running = 0.0
for xb, yb in loader:
optimizer.zero_grad()
loss = criterion(model(xb), yb)
loss.backward()
optimizer.step()
running += loss.item()
print(f"Epoch {epoch} | train {running/len(loader):.4f}")
▶ Output
Resumed from epoch 30 (last train loss 0.0232) Epoch 31 | train 0.0117 Epoch 32 | train 0.0141 Epoch 33 | train 0.0108 Epoch 34 | train 0.0115 Epoch 35 | train 0.0141
What happened here: The checkpoint is a plain dictionary holding the model’s state_dict(), the optimizer’s state_dict(), and the epoch counter. On reload we rebuild a fresh model and optimizer, pour the saved states back in, and continue from epoch 31 with loss already down near 0.02 instead of the 1.1 we started from. We pass weights_only=True to torch.load, which refuses to run any code hidden in the file and blocks a known pickle attack. In PyTorch 2.12.1 that is already the default, but writing it out states your intent clearly.
The Full Classifier End to End
Now put every piece together: the Dataset, a train and validation split, the five-step loop, both mode switches, a loss curve, and the checkpoint. This is the script you copy into a new project and adapt.
📄 train.py: the complete training and validation loop
import torch, torch.nn as nn
from torch.utils.data import DataLoader, random_split
from model_defs import ProduceDataset, FreshnessNet
torch.manual_seed(42)
ds = ProduceDataset("produce.csv")
train_ds, val_ds = random_split(ds, [len(ds) - 48, 48],
generator=torch.Generator().manual_seed(0))
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=32)
model = FreshnessNet()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
train_curve, val_curve = [], []
for epoch in range(1, 31):
model.train() # dropout ON
running = 0.0
for xb, yb in train_loader:
optimizer.zero_grad()
loss = criterion(model(xb), yb)
loss.backward()
optimizer.step()
running += loss.item()
train_loss = running / len(train_loader)
model.eval() # dropout OFF
v_loss, correct, total = 0.0, 0, 0
with torch.no_grad(): # skip the graph
for xb, yb in val_loader:
out = model(xb)
v_loss += criterion(out, yb).item()
correct += (out.argmax(1) == yb).sum().item()
total += len(yb)
train_curve.append(train_loss); val_curve.append(v_loss / len(val_loader))
if epoch == 1 or epoch % 5 == 0:
print(f"Epoch {epoch:>2d} | train {train_loss:.4f} | "
f"val {v_loss/len(val_loader):.4f} | val acc {correct/total:.1%}")
# plot the loss curves so you can see under/overfitting at a glance
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
plt.plot(range(1, 31), train_curve, label="train")
plt.plot(range(1, 31), val_curve, label="val")
plt.xlabel("epoch"); plt.ylabel("cross-entropy loss"); plt.legend()
plt.savefig("loss_curves.png", dpi=90, bbox_inches="tight")
torch.save({"epoch": 30, "model_state": model.state_dict(),
"optim_state": optimizer.state_dict(), "train_curve": train_curve},
"checkpoint.pth")
▶ Output
Epoch 1 | train 1.0987 | val 0.9213 | val acc 95.8% Epoch 5 | train 0.3787 | val 0.1755 | val acc 100.0% Epoch 10 | train 0.1365 | val 0.0467 | val acc 100.0% Epoch 15 | train 0.0533 | val 0.0130 | val acc 100.0% Epoch 20 | train 0.0512 | val 0.0054 | val acc 100.0% Epoch 25 | train 0.0263 | val 0.0034 | val acc 100.0% Epoch 30 | train 0.0232 | val 0.0022 | val acc 100.0%
What happened here: Loss falls from 1.10 to about 0.02 and validation accuracy climbs to 100 percent on this easy, well-separated dataset. The validation loss tracks the training loss closely and keeps dropping, which is the healthy shape you want. If the validation curve had turned upward while training kept falling, that gap would be overfitting, and the plot saved to loss_curves.png is exactly how you spot it at a glance. Because we seeded torch.manual_seed(42) and the split generator, your numbers will match these; drop the seeds and the shape stays the same while the exact decimals shift.
Common Mistakes
- Shuffling the validation loader: Only the training loader needs
shuffle=True. Shuffling validation changes nothing about accuracy but wastes effort and makes logs harder to compare. - Returning the wrong label dtype:
CrossEntropyLosswants class indices astorch.long. Return floats and you get a cryptic type error. - Heavy work in __init__ instead of __getitem__: For large image sets, load and transform each sample inside
__getitem__, not all at once in__init__, or you blow up memory. - Saving only the weights when you meant to resume: To continue training you also need the optimizer state and epoch, not just
model.state_dict(). - num_workers on Windows without a main guard: If you set
num_workers>0, wrap your run inif __name__ == "__main__":or the worker processes will crash.
More in this series:
- Backpropagation Explained: Gradient Descent in Python Made Simple
- Python: PyTorch vs TensorFlow vs JAX, Deep Learning Frameworks Compared
- DL: First Neural Network with TensorFlow/Keras
Frequently Asked Questions
What is the difference between a Dataset and a DataLoader in PyTorch?
A Dataset knows how to fetch one sample by index through __getitem__ and how many samples exist through __len__. A PyTorch DataLoader wraps a Dataset and handles the batching, shuffling, and parallel loading. You write the Dataset once for your data source, then reuse the same DataLoader machinery for every project. Think of the Dataset as the pantry and the DataLoader as the conveyor belt that carries batches to the model.
Should I set num_workers on the PyTorch DataLoader?
num_workers spins up extra processes that load and transform data in parallel, which speeds up training when preprocessing is the bottleneck, such as decoding and resizing images. On a CPU-only run with small tabular data it makes little difference. Start at 0 to keep things simple, then try 2 or 4 and measure. On Windows you must guard the entry point with if __name__ == ‘__main__’ or the workers fail to start.
Why shuffle the training data every epoch?
Shuffling breaks any accidental order in your file, like all the fresh samples sitting at the top. Without it the model sees batches that are not representative and the gradients get biased, which slows learning and can hurt accuracy. Set shuffle=True on the training loader and leave it off for validation and test loaders, where order does not affect the result.
Is PyTorch the only option, or are there alternatives?
PyTorch is the most widely used deep learning framework at the time of writing, but it is not the only choice. Keras 3 runs on top of PyTorch, JAX, or TensorFlow and offers a higher-level fit() API, and JAX itself is popular for research that needs aggressive compilation. The Dataset and DataLoader ideas here carry over: every framework needs a way to serve shuffled batches, so the mental model transfers even if the class names change.
Interview Questions on PyTorch DataLoader
How interviewers actually probe this topic: real scenarios, with answers you can say out loud.
Q: What two methods must a custom PyTorch Dataset implement, and what does each return?
__len__ returns the number of samples so the DataLoader knows how many batches to make, and __getitem__(idx) returns a single sample, usually a (features, label) tuple of tensors. Everything else, batching and shuffling and parallel loading, is handled by the DataLoader on top of those two methods. Keeping __getitem__ lightweight and lazy is what lets the pattern scale to datasets far larger than memory.
Q: Walk through the five steps of a single training iteration.
Call optimizer.zero_grad() to clear leftover gradients, run the forward pass out = model(x), compute loss = criterion(out, y), call loss.backward() to fill each parameter’s .grad, then optimizer.step() to nudge the weights. Miss zero_grad() and gradients accumulate across batches, producing erratic updates. That order is the same for a tiny classifier and a huge transformer.
Q: Why do you need both model.eval() and torch.no_grad() for validation? Are they the same thing?
No, they do separate jobs. model.eval() switches layers like dropout and batch normalization into inference behaviour, so predictions are deterministic and use running statistics. torch.no_grad() stops autograd from building the computation graph, which saves memory and time since you are not going to call backward. Use both: eval() for correct outputs, no_grad() for efficiency. Then switch back to model.train() before the next training pass.
Q: Scenario: to resume a long training run you saved model.state_dict() to disk, but after loading, the loss jumps back up and training feels like it restarted. What did you forget?
You saved the weights but not the optimizer state. Optimizers like Adam keep per-parameter running averages of gradients, and those got reset to zero on reload, so the first few steps overshoot until the averages rebuild. A resumable checkpoint saves optimizer.state_dict() alongside model.state_dict() and the epoch number, and reloads all three. Save the weights only when you just need to run inference later.
Q: Scenario: your input CSV has all class-0 rows first, then all class-1, then class-2, and training loss barely moves. What is the likely cause?
The DataLoader probably has shuffle=False, so early batches contain only class 0, the next only class 1, and so on. Each batch pulls the model hard toward whichever class it sees, and the updates fight each other epoch after epoch. Turning on shuffle=True mixes the classes within every batch, gradients become representative, and the loss starts falling. It is a one-word fix that people miss because the code runs without error.
What’s Next?
You now own the real PyTorch training workflow: a custom Dataset that serves samples from any source, a PyTorch DataLoader that batches and shuffles them, the five-step loop with zero_grad() in the right place, validation with both mode switches, and checkpoints you can resume from. This is the same shape behind every serious model, so you can read production training code and know exactly what each line does. Next, compare it against the higher-level style in the PyTorch vs TensorFlow comparison to see the tradeoffs between writing the loop yourself and letting a framework do it.
Want the full roadmap from Python basics through deep learning? Head back to the Python + AI/ML tutorial series home to see where this post fits and what to read next.
Reference: the complete, always-current details live in PyTorch documentation.
Related Posts
Previous: DL: First Neural Network with PyTorch
Next: DL: PyTorch Graphics Processing Unit (GPU) Training, Debugging, and Speedups
Series Home: Python + AI/ML Tutorial Series

No comment