HowToCodePage

Modules, data and saving

Writing models as nn.Module, feeding data with Dataset/DataLoader, and persisting weights correctly.

A model is a class

import torch
from torch import nn

class TabularNet(nn.Module):
    def __init__(self, n_features, n_classes):
        super().__init__()                       # always call this first
        self.net = nn.Sequential(
            nn.Linear(n_features, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, n_classes),
        )

    def forward(self, x):                        # defines the computation
        return self.net(x)

model = TabularNet(20, 3)
sum(p.numel() for p in model.parameters())
  • Submodules assigned as attributes are registered automatically — that is how .parameters() finds them.
  • Define layers in __init__, computation in forward.
  • nn.CrossEntropyLoss expects raw logits, so no softmax in the model.
  • nn.Sequential is enough for straight-line stacks; use explicit modules when you need branches or shared layers.

Dataset and DataLoader

from torch.utils.data import Dataset, DataLoader

class CSVData(Dataset):
    def __init__(self, X, y):
        self.X = torch.as_tensor(X, dtype=torch.float32)
        self.y = torch.as_tensor(y, dtype=torch.long)

    def __len__(self):
        return len(self.y)

    def __getitem__(self, i):
        return self.X[i], self.y[i]

loader = DataLoader(CSVData(X, y), batch_size=64, shuffle=True,
                    num_workers=2, pin_memory=True)
ArgumentWhy
shuffle=TrueTraining only — shuffling validation adds noise for nothing
num_workersParallel data loading; 0 on Windows if workers misbehave
pin_memorySpeeds up host-to-GPU copies
drop_lastAvoids tiny final batches that skew batchnorm

Saving and loading

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

model = TabularNet(20, 3)
model.load_state_dict(torch.load("model.pt", map_location="cpu"))
model.eval()                       # do not forget: affects dropout and batchnorm

# resuming training keeps the optimiser too
torch.save({"model": model.state_dict(),
            "optim": optimizer.state_dict(),
            "epoch": epoch}, "checkpoint.pt")
⚠️
Prefer state_dict() over pickling the whole model: it survives refactoring, is portable across devices, and does not execute arbitrary code on load. Always call model.eval() before inference.

FAQ

How do I use a pretrained model?
Load the backbone, replace the final layer for your number of classes, and optionally freeze the backbone while the new head trains. Fine-tune the whole network only when you have enough data.
Where do dropout and batchnorm bite?
They behave differently in training and evaluation. Forgetting model.eval() makes inference inconsistent and is one of the most common production bugs.

The training loop Pipelines and saving models

Last refreshed 2026-09-18.