The input pipeline

Implement __len__ and __getitem__ on a class and PyTorch already knows what to call it: a Dataset, the smallest contract in the pipeline, one item in, one item back, nothing else required. Everything that makes training practical sits one layer up, in DataLoader, which composes batching, shuffling, and prefetching on top of that bare contract, along with the option to hand the work off to separate worker processes entirely.

Loading data is concurrent Python, not tensor math.

None of this is model code, and treating it as an afterthought is the single most common way a training run ends up slower than the model itself would predict. The loader is a small concurrent program running alongside whatever device is doing the actual math, and like any concurrent program it has failure modes that have nothing to do with tensors.

Set num_workers above zero and DataLoader doesn't spin up threads inside the same process. It forks or spawns real OS processes, each running its own copy of the dataset, each handing finished batches back to the main process through a queue. That has a consequence worth stating plainly: everything a worker needs, the dataset object itself included, has to survive being pickled across that process boundary.

This is exactly where a lambda inside a dataset's transform pipeline stops working: a worker process reconstructs the dataset it was handed by unpickling it, and a lambda can't be pickled, so the transform fails before a single sample loads. The failure looks like it's about tensors or shapes; it's actually a multiprocessing constraint, and the fix is the same as it would be anywhere a lambda meets pickling: a named function, or a small callable class, in place of the anonymous one.

The reason to pay workers this way at all, instead of loading everything on the main process, is overlap. While the device is busy computing on batch n, a worker process is already reading and transforming batch n+1 off disk, so the next batch is often waiting by the time the device asks for it. It's the same overlap argument this site's kernel path makes at the level of a single Pallas grid, one level up: hide the slow transfer behind the fast compute, wherever the boundary between them happens to sit.

import torch
from torch.utils.data import DataLoader, TensorDataset

ds = TensorDataset(torch.randn(100, 8), torch.randn(100, 1))
dl = DataLoader(ds, batch_size=32, shuffle=True)
xb, yb = next(iter(dl))
print(xb.shape, yb.shape)   # torch.Size([32, 8]) torch.Size([32, 1])

Transforms, resizing an image, tokenizing a string, augmenting a sample, run inside the worker process, per sample, before the batch is ever assembled. That placement is deliberate, not an accident of the API: heavy per-sample work belongs in a worker, where it can run concurrently with the device, not inside the training step, where it would simply add to the time every step already takes.

Get this placement backward, put the heavy work in the step instead of the loader, or starve the loader with num_workers=0 and a transform that actually costs something, and the symptom looks identical from the outside: the device sits idle waiting on data that isn't ready yet. Chapter nine closes this loop with the tool that tells the two apart for certain, a profiler trace, rather than a guess about where the time went.