PyTorch nn.Module Explained: Missing Parameters, state_dict, Buffers and Registration Bugs

AI, PyTorch
Page content

PyTorch: Zero to Hero — Step 04

In the previous post we built a neural network using raw tensors and autograd.

Now we are going to add the abstraction PyTorch expects almost every real model to use:

class Model(torch.nn.Module):
    ...

But this is not going to be a tour of nn.Module methods.

The useful question for programmers is:

What exactly does nn.Module register, track, move, save and expose — and how do those mechanisms break?

If you have ever seen any of these problems:

parameter has a gradient but optimizer does not update it
parameter missing from model.parameters()
layer missing from state_dict()
model.to("cuda") moved some tensors but not others
ModuleList works but Python list does not
checkpoint loads with missing_keys / unexpected_keys
requires_grad=False but parameter still appears in state_dict

then this post is for you.

1. nn.Module is mostly a registration system

Start with a normal model:

import torch
import torch.nn as nn

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(4, 8)
        self.fc2 = nn.Linear(8, 2)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = TinyNet()

The obvious thing nn.Module gives us is callable model behavior:

x = torch.randn(16, 4)
logits = model(x)
print(logits.shape)   # torch.Size([16, 2])

But the important part is not forward().

It is registration.

When we write:

self.fc1 = nn.Linear(4, 8)

PyTorch notices that fc1 is itself an nn.Module and registers it as a submodule.
And because nn.Linear contains nn.Parameter objects, those parameters become part of the parent model’s parameter tree.

Inspect the tree:

for name, module in model.named_modules():
    print(name, '->', type(module).__name__)

Typical output:

 -> TinyNet
fc1 -> Linear
fc2 -> Linear

Now inspect parameters:

for name, param in model.named_parameters():
    print(name, param.shape)

Output:

fc1.weight torch.Size([8, 4])
fc1.bias torch.Size([8])
fc2.weight torch.Size([2, 8])
fc2.bias torch.Size([2])

We can visualise this registration hierarchy:

    graph TD
    ROOT[TinyNet] --> FC1[fc1: Linear]
    ROOT --> FC2[fc2: Linear]
    FC1 --> W1[weight: Parameter 8×4]
    FC1 --> B1[bias: Parameter 8]
    FC2 --> W2[weight: Parameter 2×8]
    FC2 --> B2[bias: Parameter 2]

This registration tree drives a huge amount of PyTorch behavior:

model.parameters()
model.named_parameters()
model.modules()
model.to(device)
model.cuda()
model.cpu()
model.state_dict()
model.train()
model.eval()

If something is not registered, many of those systems simply cannot see it.

2. A plain tensor is not a parameter

This is one of the easiest bugs to create.

class BrokenLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = torch.randn(out_features, in_features, requires_grad=True)
        self.bias = torch.zeros(out_features, requires_grad=True)

    def forward(self, x):
        return x @ self.weight.T + self.bias

The tensors can participate in autograd:

model = BrokenLinear(4, 2)
x = torch.randn(8, 4)
y = torch.randint(0, 2, (8,))

loss = nn.functional.cross_entropy(model(x), y)
loss.backward()

print(model.weight.grad is None)   # False

So gradients work.

But now:

print(list(model.parameters()))   # []

That is the trap.

Autograd knows about the tensor.
nn.Module does not know it is a model parameter.

The fix is nn.Parameter:

class ManualLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features) * 0.01)
        self.bias = nn.Parameter(torch.zeros(out_features))

    def forward(self, x):
        return x @ self.weight.T + self.bias

Now:

model = ManualLinear(4, 2)

for name, param in model.named_parameters():
    print(name, param.shape)

Output:

weight torch.Size([2, 4])
bias torch.Size([2])

The distinction is critical. We can contrast what each kind of tensor participates in:

    graph TD
    subgraph Plain Tensor
        PT[plain_tensor: requires_grad=True]
        PT --> AG[autograd ✓]
        PT --> OPT[optimizer ✗]
        PT --> SD[state_dict ✗]
        PT --> TO[model.to ✗]
    end
    subgraph nn.Parameter
        NP[nn.Parameter]
        NP --> AG2[autograd ✓]
        NP --> OPT2[optimizer ✓]
        NP --> SD2[state_dict ✓]
        NP --> TO2[model.to ✓]
    end
    subgraph Buffer
        RB[register_buffer]
        RB --> AG3[autograd ✗]
        RB --> OPT3[optimizer ✗]
        RB --> SD3[state_dict ✓*]
        RB --> TO3[model.to ✓]
    end
    style AG fill:#e6ffe6
    style AG2 fill:#e6ffe6
    style SD2 fill:#e6ffe6
    style TO2 fill:#e6ffe6
    style SD3 fill:#e6ffe6
    style TO3 fill:#e6ffe6

* – persistent buffers appear in state_dict(), non‑persistent ones do not.

3. Prove the optimizer problem

Here is a useful debugging test.

First the broken version:

model = BrokenLinear(4, 2)

try:
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
except ValueError as exc:
    print(exc)

You will get an error similar to:

optimizer got an empty parameter list

The tensor has gradients, but the model exposes no registered parameters.

Now the correct version:

model = ManualLinear(4, 2)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

print(sum(p.numel() for p in model.parameters()))   # 10

For debugging production code, I often start with:

def inspect_parameters(model):
    for name, p in model.named_parameters():
        print(
            f'{name:40s}',
            f'shape={tuple(p.shape)!s:18s}',
            f'requires_grad={str(p.requires_grad):5s}',
            f'grad={"None" if p.grad is None else "set"}'
        )

Then:

inspect_parameters(model)

If the parameter you expect is missing from this output, your training loop is not the first place to debug.

Registration is.

4. The Python list bug

This one catches experienced Python programmers because ordinary Python containers look perfectly reasonable.

Consider this dynamic network:

class BrokenMLP(nn.Module):
    def __init__(self, width=32, depth=4):
        super().__init__()
        self.layers = [
            nn.Linear(width, width)
            for _ in range(depth)
        ]

    def forward(self, x):
        for layer in self.layers:
            x = torch.relu(layer(x))
        return x

The forward pass works:

model = BrokenMLP()
x = torch.randn(8, 32)
print(model(x).shape)   # torch.Size([8, 32])

But inspect the parameter tree:

print(sum(p.numel() for p in model.parameters()))   # 0

Why?

Because self.layers is a normal Python list.
nn.Module cannot recursively discover arbitrary objects buried inside arbitrary Python containers.

Use nn.ModuleList:

class MLP(nn.Module):
    def __init__(self, width=32, depth=4):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(width, width)
            for _ in range(depth)
        ])

    def forward(self, x):
        for layer in self.layers:
            x = torch.relu(layer(x))
        return x

Now:

model = MLP()
print(sum(p.numel() for p in model.parameters()))   # 4224

And the names are visible:

for name, _ in model.named_parameters():
    print(name)

Example:

layers.0.weight
layers.0.bias
layers.1.weight
layers.1.bias
layers.2.weight
layers.2.bias
layers.3.weight
layers.3.bias

We can visualise the difference:

    graph TD
    subgraph Python list
        PL[model.layers = Python list] --> M1[Layer 0]
        PL --> M2[Layer 1]
        PL --> M3[Layer ...]
        PL -.-> |Not registered| REG[model.parameters = empty]
    end
    subgraph nn.ModuleList
        ML[model.layers = ModuleList] --> M4[Layer 0]
        ML --> M5[Layer 1]
        ML --> M6[Layer ...]
        ML --> |Registered| REG2[model.parameters contains all]
    end
    style REG fill:#ffcccc
    style REG2 fill:#ccffcc

The same rule applies to dynamic parameter collections. Use:

nn.ParameterList
nn.ParameterDict
nn.ModuleList
nn.ModuleDict

not plain containers when you need PyTorch registration.

5. state_dict() is the model’s registered state

Inspect it:

model = TinyNet()
state = model.state_dict()

for key, value in state.items():
    print(key, tuple(value.shape))

Output:

fc1.weight (8, 4)
fc1.bias (8,)
fc2.weight (2, 8)
fc2.bias (2,)

The names follow the module hierarchy.

Think of them as paths:

fc1.weight
│   └── parameter
└────── submodule

Nested modules produce nested names.

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.block = nn.Sequential(
            nn.Linear(4, 8),
            nn.ReLU(),
            nn.Linear(8, 8),
        )

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.head = nn.Linear(8, 2)

model = Model()

for key in model.state_dict():
    print(key)

Output:

encoder.block.0.weight
encoder.block.0.bias
encoder.block.2.weight
encoder.block.2.bias
head.weight
head.bias

This naming is why changing model structure often changes checkpoint compatibility.

6. state_dict() is not just trainable parameters

A common misconception is:

state_dict == model.parameters()

Not quite.

state_dict() contains registered parameters and persistent buffers.

Buffers are model state that should move with the model and often be saved, but should not be optimized.

Example:

class RunningMean(nn.Module):
    def __init__(self, features):
        super().__init__()
        self.register_buffer('running_mean', torch.zeros(features))

    def forward(self, x):
        return x - self.running_mean

Inspect it:

model = RunningMean(4)

print('parameters:', list(model.named_parameters()))
print('buffers:', list(model.named_buffers()))
print('state_dict:', list(model.state_dict().keys()))

Output:

parameters: []
buffers: [('running_mean', tensor([0., 0., 0., 0.]))]
state_dict: ['running_mean']

That makes buffers ideal for things such as:

running statistics
normalization constants
masks
lookup state
non-learnable tensors tied to the model

7. Why not just assign a normal tensor?

Compare:

class BadState(nn.Module):
    def __init__(self):
        super().__init__()
        self.scale = torch.ones(4)

with:

class GoodState(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer('scale', torch.ones(4))

Now inspect:

bad = BadState()
good = GoodState()

print(bad.state_dict().keys())   # odict_keys([])
print(good.state_dict().keys())  # odict_keys(['scale'])

And device movement differs too.

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

good = good.to(device)
print(good.scale.device)   # the buffer moved

The registered buffer follows the module.
A random attribute containing a tensor is not part of that machinery.

8. Persistent and non-persistent buffers

Sometimes a tensor should move with the model but should not be written into checkpoints.

PyTorch supports that directly:

class PositionalCache(nn.Module):
    def __init__(self, size):
        super().__init__()
        self.register_buffer(
            'cache',
            torch.arange(size),
            persistent=False,
        )

Now:

model = PositionalCache(16)

print(list(model.named_buffers()))          # cache is present
print(list(model.state_dict().keys()))      # cache is absent

The buffer exists and moves with the module, but is absent from state_dict().

This is useful for caches that can be recomputed.

9. requires_grad=False does not unregister a parameter

Consider:

model = TinyNet()

model.fc1.weight.requires_grad_(False)
model.fc1.bias.requires_grad_(False)

Now:

for name, p in model.named_parameters():
    print(name, p.requires_grad)

The frozen parameters are still present.
And they are still in state_dict():

print(model.state_dict().keys())

Freezing means:

do not calculate gradients for this parameter

It does not mean:

remove this parameter from the model

That distinction matters when fine-tuning.

A common optimizer pattern is:

trainable = [
    p for p in model.parameters()
    if p.requires_grad
]

optimizer = torch.optim.AdamW(trainable, lr=1e-3)

10. Debug what the optimizer actually owns

The model and optimizer are separate objects.

A parameter can be registered in the model and still not belong to the optimizer.

Use this helper:

def optimizer_parameter_ids(optimizer):
    return {
        id(p)
        for group in optimizer.param_groups
        for p in group['params']
    }


def audit_optimizer(model, optimizer):
    owned = optimizer_parameter_ids(optimizer)

    for name, p in model.named_parameters():
        print(
            f'{name:40s}',
            'trainable=' + str(p.requires_grad),
            'in_optimizer=' + str(id(p) in owned),
        )

Example:

model = TinyNet()
optimizer = torch.optim.AdamW(model.fc2.parameters(), lr=1e-3)

audit_optimizer(model, optimizer)

Output will show that fc1 parameters are registered and trainable but not owned by the optimizer.

That is a very different bug from missing gradients.

11. Replacing parameters after creating the optimizer

This is subtle and dangerous.

model = TinyNet()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

Now replace a layer:

model.fc2 = nn.Linear(8, 2)

The model now contains a new fc2.
But the optimizer was constructed using references to the old parameters.

Check it:

audit_optimizer(model, optimizer)

The new fc2 parameters can be missing from the optimizer.

Fix:

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

General rule:

If you replace registered parameters or modules after optimizer construction, verify optimizer membership.

12. model.to(device) only moves registered state

Suppose we write:

class DeviceBug(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(4, 4))
        self.mask = torch.ones(4, 4)

    def forward(self, x):
        return (x @ self.weight) * self.mask

Move it:

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

print('weight:', model.weight.device)   # cuda:0 (or cpu)
print('mask:', model.mask.device)       # cpu -> mismatch!

On CUDA you would see the bug: the mask stayed on CPU.

Fix the model state definition:

class DeviceSafe(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(4, 4))
        self.register_buffer('mask', torch.ones(4, 4))

    def forward(self, x):
        return (x @ self.weight) * self.mask

Now both tensors move together.

13. A device audit helper

For complicated models:

def audit_devices(model):
    print('PARAMETERS')
    for name, p in model.named_parameters():
        print(f'{name:50s} {p.device}')

    print('\nBUFFERS')
    for name, b in model.named_buffers():
        print(f'{name:50s} {b.device}')

And for input batches:

def assert_same_device(model, *tensors):
    devices = {
        p.device
        for p in model.parameters()
    }

    devices.update(t.device for t in tensors)

    if len(devices) != 1:
        raise RuntimeError(f'device mismatch: {devices}')

This turns a vague runtime failure into a deliberate contract.

14. ModuleDict for named dynamic modules

Sometimes names matter.

class MultiHeadModel(nn.Module):
    def __init__(self, width):
        super().__init__()

        self.heads = nn.ModuleDict({
            'sentiment': nn.Linear(width, 3),
            'topic': nn.Linear(width, 20),
            'spam': nn.Linear(width, 2),
        })

    def forward(self, x, task):
        return self.heads[task](x)

Inspect:

model = MultiHeadModel(128)

for name, p in model.named_parameters():
    print(name)

Output:

heads.sentiment.weight
heads.sentiment.bias
heads.topic.weight
heads.topic.bias
heads.spam.weight
heads.spam.bias

This is much safer than hiding modules in an ordinary dictionary.

15. ParameterList and ParameterDict

If you need raw trainable tensors rather than submodules:

class LearnedVectors(nn.Module):
    def __init__(self, count, width):
        super().__init__()
        self.vectors = nn.ParameterList([
            nn.Parameter(torch.randn(width))
            for _ in range(count)
        ])

Or named parameters:

class Gates(nn.Module):
    def __init__(self):
        super().__init__()
        self.gates = nn.ParameterDict({
            'input': nn.Parameter(torch.tensor(1.0)),
            'output': nn.Parameter(torch.tensor(1.0)),
        })

The rule is consistent:

ModuleList / ModuleDict       -> collections of modules
ParameterList / ParameterDict -> collections of parameters

16. Saving a model correctly

The usual pattern is:

torch.save(model.state_dict(), 'model.pt')

Then recreate the architecture:

model = TinyNet()
state = torch.load('model.pt', map_location='cpu')
model.load_state_dict(state)

The checkpoint contains state.
Your Python code defines architecture.

That separation is valuable.

Inspect a checkpoint before loading:

state = torch.load('model.pt', map_location='cpu')

for key, tensor in state.items():
    print(key, tensor.shape, tensor.dtype)

When debugging model loading, inspect the keys before doing anything clever.

17. strict=True is your friend

Imagine version 1:

class ModelV1(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Linear(4, 8)
        self.head = nn.Linear(8, 2)

Then version 2:

class ModelV2(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Linear(4, 8)
        self.classifier = nn.Linear(8, 2)

We renamed head to classifier.

Loading the old checkpoint normally:

v2 = ModelV2()
state = ModelV1().state_dict()

v2.load_state_dict(state)

will fail because the keys do not match.

That is good. The error tells you the architectural contract changed.

You can inspect explicitly:

result = v2.load_state_dict(state, strict=False)

print('missing:', result.missing_keys)
print('unexpected:', result.unexpected_keys)

Typical result:

missing: ['classifier.weight', 'classifier.bias']
unexpected: ['head.weight', 'head.bias']

Do not use strict=False merely to silence errors.
Use it when you understand why the key sets differ.

18. Remap checkpoint keys deliberately

If the rename is intentional:

old_state = ModelV1().state_dict()

new_state = {}

for key, value in old_state.items():
    if key.startswith('head.'):
        key = key.replace('head.', 'classifier.', 1)

    new_state[key] = value

model = ModelV2()
result = model.load_state_dict(new_state)

For larger migrations, write a transformation function and test it.

def migrate_v1_to_v2(state):
    migrated = {}

    for key, value in state.items():
        if key.startswith('head.'):
            key = key.replace('head.', 'classifier.', 1)

        migrated[key] = value

    return migrated

Then assert:

migrated = migrate_v1_to_v2(old_state)
model.load_state_dict(migrated, strict=True)

Checkpoint migration is code.
Treat it like code.

19. Compare model schemas

A very useful utility:

def state_schema(model):
    return {
        name: {
            'shape': tuple(t.shape),
            'dtype': str(t.dtype),
        }
        for name, t in model.state_dict().items()
    }

Compare two versions:

v1 = state_schema(ModelV1())
v2 = state_schema(ModelV2())

only_v1 = v1.keys() - v2.keys()
only_v2 = v2.keys() - v1.keys()
common = v1.keys() & v2.keys()

print('only v1:', sorted(only_v1))
print('only v2:', sorted(only_v2))

for key in sorted(common):
    if v1[key] != v2[key]:
        print('changed:', key, v1[key], '->', v2[key])

This is far more useful than staring at a giant load error.

20. Detect parameters that never receive gradients

After backward:

loss.backward()

run:

def find_missing_gradients(model):
    missing = []

    for name, p in model.named_parameters():
        if p.requires_grad and p.grad is None:
            missing.append(name)

    return missing

Usage:

missing = find_missing_gradients(model)
print(missing)

A parameter can be correctly registered and correctly owned by the optimizer but still receive no gradient because the forward graph never used it.

Again, separate the layers of debugging. We can formalise this as a decision flow:

    flowchart TD
    A[Parameter] --> B{Registered?}
    B -- No --> B1[Add as nn.Parameter or register_buffer]
    B -- Yes --> C{requires_grad?}
    C -- No --> C1[Maybe intentionally frozen]
    C -- Yes --> D{Used in forward graph?}
    D -- No --> D1[Check forward code]
    D -- Yes --> E{Gradient produced?}
    E -- No --> E1[Check graph connectivity]
    E -- Yes --> F{In optimizer?}
    F -- No --> F1[Add to param_groups]
    F -- Yes --> G{Value changed after step?}
    G -- No --> G1[Check LR, zero_grad, etc.]
    G -- Yes --> H[All good]

21. Prove parameters actually changed

A robust helper:

def snapshot_parameters(model):
    return {
        name: p.detach().clone()
        for name, p in model.named_parameters()
    }


def changed_parameters(model, before):
    changed = []

    for name, p in model.named_parameters():
        if not torch.equal(p.detach(), before[name]):
            changed.append(name)

    return changed

Use it around a training step:

before = snapshot_parameters(model)

optimizer.zero_grad()
loss = nn.functional.cross_entropy(model(x), y)
loss.backward()
optimizer.step()

print(changed_parameters(model, before))

This is the final test.

Not:

loss.backward() ran

Not:

p.grad is not None

But:

the parameters actually changed

22. A complete model audit

Here is a reusable debugging utility for real projects:

def audit_model(model, optimizer=None):
    print('=== MODULES ===')
    for name, module in model.named_modules():
        print(name or '<root>', type(module).__name__)

    print('\n=== PARAMETERS ===')
    optimizer_ids = None

    if optimizer is not None:
        optimizer_ids = {
            id(p)
            for group in optimizer.param_groups
            for p in group['params']
        }

    for name, p in model.named_parameters():
        print(
            name,
            'shape=', tuple(p.shape),
            'device=', p.device,
            'requires_grad=', p.requires_grad,
            'grad=', p.grad is not None,
            'optimizer=', (
                id(p) in optimizer_ids
                if optimizer_ids is not None
                else 'n/a'
            ),
        )

    print('\n=== BUFFERS ===')
    for name, b in model.named_buffers():
        print(
            name,
            'shape=', tuple(b.shape),
            'device=', b.device,
        )

    print('\n=== STATE DICT ===')
    for name, tensor in model.state_dict().items():
        print(name, tuple(tensor.shape), tensor.dtype)

This one helper can answer a surprising number of PyTorch questions.

23. Why nn.Module matters

After the previous post, we already know how to build a network without it.

So why use it?

Because a real model needs a coherent state graph.

nn.Module gives us that graph.

It turns this:

some Python objects containing tensors

into this:

registered model
├── submodules
├── parameters
├── buffers
└── named persistent state

That structure allows PyTorch to implement:

optimization discovery
checkpointing
device transfer
nested models
freezing
introspection
training/evaluation mode propagation

The abstraction is not primarily there to make forward() look pretty.

It is there to make the model’s state explicit and traversable.

24. The registration checklist

When something strange happens, check these in order.

Is the layer registered?

print(dict(model.named_modules()).keys())

Is the parameter registered?

print(dict(model.named_parameters()).keys())

Is persistent state registered?

print(dict(model.named_buffers()).keys())

Is it in the checkpoint?

print(model.state_dict().keys())

Is it trainable?

print(param.requires_grad)

Did it receive a gradient?

print(param.grad)

Does the optimizer own it?

id(param) in optimizer_parameter_ids(optimizer)

Did it actually change?

torch.equal(before, param.detach())

That progression is much faster than randomly rewriting a training loop.

25. Challenge: break registration deliberately

Build this model:

class ExperimentalNet(nn.Module):
    def __init__(self, width=32):
        super().__init__()

        self.input = nn.Linear(width, width)

        self.blocks = [
            nn.Linear(width, width)
            for _ in range(3)
        ]

        self.scale = torch.ones(width)

        self.output = nn.Linear(width, 2)

    def forward(self, x):
        x = torch.relu(self.input(x))

        for block in self.blocks:
            x = torch.relu(block(x))

        x = x * self.scale
        return self.output(x)

Then answer these without guessing:

Which weights appear in model.parameters()?
Which tensors appear in state_dict()?
Which tensors move with model.to(device)?
Which tensors receive gradients?
Which tensors can an optimizer built from model.parameters() update?

Then repair it with:

nn.ModuleList
register_buffer

Run audit_model() before and after.

The goal is to see the registration system directly.

Where the series goes next

We have now moved from raw tensors to a real model abstraction:

Step 00 — What Are We Actually Doing?
Step 01 — Tensor Shapes and Broadcasting Bugs
Step 02 — Autograd Debugging
Step 03 — Build a Neural Network Without nn.Module
Step 04 — nn.Module, Parameters, Buffers and state_dict

The next step is where models stop training on one giant tensor and start consuming real datasets.

And that opens another large class of bugs programmers search for constantly:

PyTorch DataLoader slow
num_workers makes training slower
CUDA GPU waiting for data
pin_memory does nothing
workers hang on Windows
DataLoader duplicates data
training stalls between batches

So Step 05 will be:

PyTorch DataLoader Performance: num_workers, pin_memory, Prefetching and Why Your GPU Is Waiting

That will introduce Dataset, DataLoader, batching and train/validation splits through the performance problems that make them matter in real systems.