PyTorch · Class

DataLoader

DataLoader wraps a dataset in an iterable that can customize sample order, batch samples, use single or multiple processes, and optionally pin fetched tensors in memory.

How it works

DataLoader wraps a map-style or iterable-style dataset and exposes it as an iterable. It can select the loading order, form batches through collation, load in one or more processes, and place fetched tensors in pinned memory for transfers to CUDA-enabled GPUs.

Syntax

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0, collate_fn=None, pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None, multiprocessing_context=None, generator=None, *, prefetch_factor=None, persistent_workers=False, pin_memory_device='', in_order=True)

Parameters

Name	Type	Description
`dataset`	`Dataset`	Dataset from which to load data.
`batch_size`	`int, optional`	Number of samples per batch; default is 1.
`shuffle`	`bool, optional`	Whether to reshuffle data at every epoch; default is False.
`sampler`	`Sampler or Iterable, optional`	Strategy used to draw samples from the dataset.
`batch_sampler`	`Sampler or Iterable, optional`	Sampler that yields mini-batches of indices.
`num_workers`	`int, optional`	Number of worker processes used for loading; default is 0.
`collate_fn`	`callable, optional`	Function used to merge samples into a batch.
`pin_memory`	`bool, optional`	Whether to copy fetched tensors into pinned memory; default is False.
`drop_last`	`bool, optional`	Whether to discard the final incomplete batch; default is False.
`timeout`	`numeric, optional`	Maximum waiting time for a batch from workers; default is 0.
`worker_init_fn`	`callable, optional`	Initialization function run for each worker.
`multiprocessing_context`	`str or multiprocessing context, optional`	Context used to start worker processes.

Returns

An iterable over the given dataset that yields individual samples or automatically collated batches, according to the configured options.

Key points

DataLoader is consumed directly in a loop, yielding data from the wrapped dataset.
The class works with both map-style datasets and iterable-style datasets.
Sampling configuration lets callers control the order in which dataset elements are fetched.
Loading may run in the main process or use multiple worker processes.
The pin_memory option can automatically place fetched tensors in pinned memory.

Common mistakes

Parallel workers are not enabled unless num_workers is set to a positive integer.
Increasing num_workers changes loading to multi-process execution, so worker-related dataset behavior must be considered.
pin_memory is useful for faster CUDA transfers, but it changes where fetched tensors are placed.

Example: Batch and shuffle a dataset

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

x = torch.arange(10).reshape(5, 2)
y = torch.arange(5)
loader = DataLoader(TensorDataset(x, y), batch_size=2, shuffle=True)

for features, labels in loader:
    print(features.shape, labels.shape)

Creates a dataset of five feature-label pairs, then iterates over shuffled mini-batches of two samples.

When to use it

Produce mini-batches from datasets

Training loops

Reshuffle samples between epochs

Randomized training

Parallelize data loading with worker processes

Input pipelines

Pin fetched tensors for faster GPU transfer

CUDA training

Interview questions

QWhat does DataLoader provide?

DataLoader is a Python iterable that combines a dataset with a sampling strategy. It supports map-style and iterable-style datasets, batching, custom loading order, single- or multi-process loading, and optional memory pinning.

Follow-up: How do you enable parallel loading? Set num_workers to a positive integer; zero retains the default single-process behavior.

QHow are batches formed?

batch_size controls how many samples are grouped into each batch. Automatic collation combines individual fetched samples, while collate_fn can customize that combination.

Follow-up: Which other arguments affect automatic batching? drop_last, batch_sampler, and collate_fn also participate in batching behavior.

QHow does DataLoader control sample order?

The dataset is the main constructor input and identifies the data source. shuffle can request reshuffling at every epoch, while sampler defines the strategy for drawing samples.

Follow-up: Can shuffle and sampler be freely combined? No: when a sampler is supplied, it determines the sampling strategy, so the two options should not be configured incompatibly.

torch.utils.data.Dataset
torch.utils.data.IterableDataset
torch.utils.data.Sampler
torch.utils.data.BatchSampler
torch.utils.data.TensorDataset
torch.utils.data.default_collate

Source: torch.utils.data — PyTorch main documentation. Checked 2026-10-15.