Introduction

When custom labels do not appear correctly in a PyTorch DataLoader, the issue is usually in dataset return format, label dtype, or collate behavior. DataLoader itself is mostly a batching and sampling wrapper, so label bugs usually originate in Dataset.__getitem__. This guide shows a robust pattern for custom labels and common debugging steps.

Implement Dataset Return Contract Correctly

A custom dataset should return one sample per index, usually as (input, label) or a dictionary with stable keys.

1import torch
2from torch.utils.data import Dataset
3
4
5class ToyDataset(Dataset):
6    def __init__(self, texts, label_names):
7        self.texts = texts
8        classes = sorted(set(label_names))
9        self.label_to_id = {name: i for i, name in enumerate(classes)}
10        self.labels = [self.label_to_id[name] for name in label_names]
11
12    def __len__(self):
13        return len(self.texts)
14
15    def __getitem__(self, idx):
16        x = torch.tensor([len(self.texts[idx])], dtype=torch.float32)
17        y = torch.tensor(self.labels[idx], dtype=torch.long)
18        return x, y

Key points:

• labels used with CrossEntropyLoss should be integer class IDs,
• label tensor dtype should be torch.long,
• output structure must be consistent for every sample.

Batch with DataLoader and Validate Shapes

Create loader and inspect one batch early.

1from torch.utils.data import DataLoader
2
3texts = ["cat", "elephant", "dog", "whale"]
4label_names = ["animal", "animal", "animal", "animal"]
5
6ds = ToyDataset(texts, label_names)
7dl = DataLoader(ds, batch_size=2, shuffle=True)
8
9xb, yb = next(iter(dl))
10print(xb.shape, xb.dtype)
11print(yb.shape, yb.dtype)
12print(yb)

If labels are missing or malformed, stop and inspect __getitem__ output directly:

sample_x, sample_y = ds[0]
print(sample_x, sample_y)

This isolates dataset logic from loader logic.

Handle Complex Labels with a Custom collate_fn

For multi-label targets, variable-length labels, or metadata-rich samples, default collation may fail. Use a custom collate function.

1from torch.nn.utils.rnn import pad_sequence
2
3
4def collate_fn(batch):
5    xs, ys = zip(*batch)
6    xs = torch.stack(xs)
7    ys = [torch.tensor(y, dtype=torch.long) for y in ys]
8    ys = pad_sequence(ys, batch_first=True, padding_value=-1)
9    return xs, ys

Then pass it to loader:

dl = DataLoader(ds, batch_size=4, collate_fn=collate_fn)

Use this when label structure is not a fixed scalar per sample.

Multi-Worker and Transform Caveats

If labels look correct with num_workers=0 but break with workers enabled, the dataset may rely on non-picklable state or mutable globals.

Debug sequence:

1. run with num_workers=0,
2. print sample types and dtypes,
3. increase workers after correctness is confirmed.

Also verify transforms do not discard labels. Some pipelines transform only image tensors and accidentally return transformed input without the label part.

End-to-End Training Example

A minimal train step confirms labels integrate with loss correctly.

1import torch
2import torch.nn as nn
3
4model = nn.Linear(1, len(ds.label_to_id))
5criterion = nn.CrossEntropyLoss()
6optim = torch.optim.Adam(model.parameters(), lr=1e-3)
7
8for xb, yb in dl:
9    logits = model(xb)
10    loss = criterion(logits, yb)
11    optim.zero_grad()
12    loss.backward()
13    optim.step()
14    print("loss", float(loss))
15    break

If this fails, inspect class count, label range, and label dtype first.

Keep Label Encoding Stable Across Runs

If labels are string classes, persist the mapping so training, validation, and inference use identical IDs.

1import json
2
3mapping_path = "label_to_id.json"
4with open(mapping_path, "w", encoding="utf-8") as f:
5    json.dump(ds.label_to_id, f, indent=2)
6
7with open(mapping_path, "r", encoding="utf-8") as f:
8    restored = json.load(f)
9
10print(restored)

Stable mappings prevent silent class swaps when dataset ordering changes between runs.

Common Pitfalls

A common mistake is returning string labels directly while using CrossEntropyLoss. Convert labels to integer class IDs before batching.

Another issue is inconsistent return shapes from __getitem__, which causes collation errors that appear like loader failures.

Developers also forget that labels for classification losses need torch.long and valid class index range. Float labels or out-of-range IDs cause runtime errors or wrong training behavior.

Summary

• Label issues usually come from dataset output, not DataLoader internals.
• Return stable (input, label) structures with correct label dtype.
• Encode class labels to integer IDs for classification losses.
• Use collate_fn for complex label shapes.
• Validate one sample and one batch before starting full training.