Introduction

In PyTorch, k-fold cross-validation is usually built by splitting dataset indices fold by fold and then creating new DataLoader objects for the training and validation subsets. The important details are to rebuild the model for each fold, keep the data split logic separate from the DataLoader, and avoid leaking weights or preprocessing state between folds.

The Basic Idea

The workflow is:

1. create one dataset
2. split its indices with a cross-validation splitter
3. wrap each fold in Subset
4. build fold-specific DataLoader objects
5. train and evaluate a fresh model on each fold

DataLoader does not perform the fold split by itself. It only loads batches from whatever dataset or subset you give it.

A Minimal End-to-End Example

Here is a small example using TensorDataset and scikit-learn’s KFold:

1import numpy as np
2import torch
3from sklearn.model_selection import KFold
4from torch import nn
5from torch.utils.data import DataLoader, Subset, TensorDataset
6
7X = torch.randn(200, 10)
8y = (X.sum(dim=1) > 0).float().unsqueeze(1)
9
10dataset = TensorDataset(X, y)
11
12def build_model():
13    return nn.Sequential(
14        nn.Linear(10, 16),
15        nn.ReLU(),
16        nn.Linear(16, 1),
17        nn.Sigmoid(),
18    )
19
20def train_one_epoch(model, loader, optimizer, loss_fn):
21    model.train()
22    total_loss = 0.0
23
24    for xb, yb in loader:
25        optimizer.zero_grad()
26        preds = model(xb)
27        loss = loss_fn(preds, yb)
28        loss.backward()
29        optimizer.step()
30        total_loss += loss.item() * xb.size(0)
31
32    return total_loss / len(loader.dataset)
33
34def evaluate(model, loader):
35    model.eval()
36    correct = 0
37
38    with torch.no_grad():
39        for xb, yb in loader:
40            preds = (model(xb) >= 0.5).float()
41            correct += (preds == yb).sum().item()
42
43    return correct / len(loader.dataset)
44
45kfold = KFold(n_splits=5, shuffle=True, random_state=42)
46fold_scores = []
47
48for fold, (train_idx, val_idx) in enumerate(kfold.split(range(len(dataset))), start=1):
49    train_subset = Subset(dataset, train_idx)
50    val_subset = Subset(dataset, val_idx)
51
52    train_loader = DataLoader(train_subset, batch_size=32, shuffle=True)
53    val_loader = DataLoader(val_subset, batch_size=32, shuffle=False)
54
55    model = build_model()
56    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
57    loss_fn = nn.BCELoss()
58
59    for epoch in range(5):
60        train_loss = train_one_epoch(model, train_loader, optimizer, loss_fn)
61
62    val_score = evaluate(model, val_loader)
63    fold_scores.append(val_score)
64    print(f"Fold {fold}: accuracy={val_score:.4f}")
65
66print("Mean accuracy:", float(np.mean(fold_scores)))

This pattern is the standard PyTorch answer: the cross-validation logic manages indices, and the DataLoader handles batching.

Rebuild the Model on Every Fold

One rule matters more than almost anything else: do not reuse the same model instance across folds.

Every fold must start from a fresh model and a fresh optimizer. Otherwise:

• weights leak from one fold to the next
• optimizer state leaks too
• the validation scores stop meaning what cross-validation is supposed to mean

That is why the example calls build_model() inside the fold loop.

Use the Right Splitter for the Problem

KFold is a good default for generic tabular regression or balanced data, but not every dataset should use it.

You may need:

• 'StratifiedKFold for imbalanced classification'
• grouped splits when related samples must stay together
• time-aware splitting for sequential data

The DataLoader code barely changes, but the fold splitter absolutely matters for honest evaluation.

Keep Transforms and Preprocessing Honest

If your dataset applies normalization, tokenization, or augmentation, make sure fold boundaries are respected. For example, statistics used for normalization should come from the training fold, not from the full dataset.

The easiest mistake is to compute preprocessing once on all data and then cross-validate afterward. That can leak validation information into training.

The DataLoader will not prevent that for you. It only loads what you tell it to load.

When Samplers Are Useful

You can also implement fold logic with samplers instead of Subset. For many projects, Subset is simpler and easier to read. Samplers become useful when:

• the dataset should stay whole
• you need custom sampling behavior
• you want tighter control over index ordering

But the conceptual model is the same: the split is index-driven.

Common Pitfalls

• Reusing the same model or optimizer across folds.
• Letting preprocessing statistics come from the full dataset instead of the training fold.
• Using plain KFold on a problem that really needs stratified or grouped splitting.
• Expecting DataLoader to perform cross-validation splitting automatically.
• Comparing fold scores without keeping training epochs and hyperparameters consistent across folds.

Summary

• In PyTorch, k-fold cross-validation is usually built by splitting dataset indices and creating fold-specific DataLoader objects.
• 'Subset is the simplest way to turn fold indices into train and validation datasets.'
• Rebuild the model and optimizer on every fold so state does not leak.
• Choose the splitter that matches the data structure, not just the most common one.
• Keep preprocessing honest, because DataLoader helps with batching, not with evaluation design.