Create a DataLoader

easyPRODatasetsPyTorch DatasetsPytorch

Problem Statement

Create a DataLoader to batch a dataset.

Background

DataLoader wraps a Dataset and provides automatic batching, shuffling, and parallel loading — the standard way to feed data into a training loop.

Your Task

Write a function create_dataloader(dataset, batch_size=2, shuffle=False) that returns a DataLoader for the given dataset with the specified batch size, shuffling only when shuffle is true (default: no shuffling).

Output Format

Return a DataLoader object.

Example

Input

dataset (6 samples), batch_size=2, shuffle=False

Output

DataLoader with 3 batches

Reasoning

Input: dataset holds $N = 6$ samples, each a (feature vector, label) pair, e.g. features $[1, 2], [3, 4], \ldots, [11, 12]$ with labels $0, 1, 0, 1, 0, 1$ . batch_size $= 2$ is how many samples get grouped together, and shuffle is off, so samples keep their original order.
Math: the number of batches is $\lceil N / B \rceil = \lceil 6 / 2 \rceil = 3$ . Batch 1 is samples 1–2, batch 2 is samples 3–4, batch 3 is samples 5–6.
What one batch looks like: the samples are stacked along a new leading axis, so batch 1 gives features $\begin{bmatrix}1 & 2\\ 3 & 4\end{bmatrix}$ (a $2 \times 2$ array: batch size $\times$ feature count) and labels $[0, 1]$ (length 2).
Output: an iterable object, not a list of numbers; its length is the batch count, 3.
Notice that the last batch can be smaller: with $N = 10$ and $B = 3$ you get $\lceil 10/3 \rceil = 4$ batches of sizes $3, 3, 3, 1$ .

Constraints

Use torch.utils.data.DataLoader
batch_size will be a positive integer
shuffle will be a boolean

Learn this in

Python FoundationsFoundations PyTorch FundamentalsFoundations

More from Pytorch

Apply Softmax for Predictionseasy Count Model Parametersmedium Build with nn.Sequentialeasy Use nn.Flatten Layereasy Create a Simple Neural Network Classeasy

Concept primerPREVIEW

1. Background Knowledge

PyTorch Data Pipeline Fundamentals PyTorch training relies on a Dataset → DataLoader → Model pipeline. The Dataset class defines how to access individual samples via getitem(idx) and len(). The DataLoader automates batching, shuffling, and parallel loading.

Key Concepts:

Batching: Stack N samples into tensors of shape (batch_size, *sample_shape)
Collation: default_collate() automatically stacks tensors, lists, or dicts
Iterators: DataLoader is iterable; next(iter(loader)) yields first batch
No shuffling: shuffle=False preserves dataset order (important for testing)

Mathematical Perspective: For dataset size $N$ and batch size $B$ , DataLoader yields $\lceil N/B \rceil$ batches. Last batch may be smaller unless drop_last=True.

2. Algorithm Approach

Core Algorithm: Wrapper pattern

Input: Dataset D, batch_size B
Output: DataLoader instance
1. Initialize DataLoader(D, batch_size=B, shuffle=False, num_workers=0)
2. Return loader

DataLoader Parameters (most relevant):

Parameter	Purpose	Default	This Problem
batch_size	Samples per batch	1	Required
shuffle	Randomize order	False	False
num_workers	Parallel loading	0	0 (simple)
drop_last	Drop incomplete batch	False	False

Collation Flow:

[sample1, sample2,...] → default_collate() → (features_batch, labels_batch)

3. Step-by-Step Strategy

from torch.utils.data import DataLoader

def create_dataloader(dataset, batch_size):
 # Step 1: Create DataLoader with exact parameters
 loader = DataLoader(
 dataset,
 batch_size=batch_size, # Required parameter
 shuffle=False # No shuffling per problem spec
 )
 return loader

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results