How to use a DataLoader in PyTorch

A PyTorch DataLoader turns a dataset into the mini-batches that training and evaluation loops consume. It controls how many samples arrive together, whether sample indexes are shuffled, and what tensor shapes the model code receives from each iteration.

For map-style datasets such as TensorDataset or a custom Dataset class, DataLoader requests dataset items by index and collates them into one batch. A dataset that returns feature and label pairs yields one feature tensor and one label tensor per loader iteration, with the batch dimension first.

Start with a small dataset, num_workers=0, and a fixed torch.Generator seed while checking the batch contract. After the batch shape, label order, and loop unpacking match the model code, the same loader pattern can move into the real training loop.

1. Create a minimal DataLoader smoke script.

   dataloader_batch_demo.py

   import torch
   from torch.utils.data import DataLoader, TensorDataset
    
    
   features = torch.arange(24, dtype=torch.float32).reshape(6, 4)
   labels = torch.arange(6)
   dataset = TensorDataset(features, labels)
    
   loader = DataLoader(
       dataset,
       batch_size=2,
       shuffle=True,
       generator=torch.Generator().manual_seed(17),
   )
    
   batch_features, batch_labels = next(iter(loader))
    
   print(f"dataset length: {len(dataset)}")
   print(f"batch features shape: {tuple(batch_features.shape)}")
   print(f"batch labels: {batch_labels.tolist()}")
   print("batch features:")
   print(batch_features)
    
   for epoch in range(2):
       epoch_loader = DataLoader(
           dataset,
           batch_size=2,
           shuffle=True,
           generator=torch.Generator().manual_seed(17),
       )
       label_order = []
       for _, batch_targets in epoch_loader:
           label_order.extend(batch_targets.tolist())
       print(f"epoch {epoch} label order: {label_order}")

   batch_size=2 asks for two samples per iteration. shuffle=True randomizes dataset indexes, and the seeded torch.Generator keeps the smoke output repeatable.

2. Run the DataLoader smoke script.

   $ python dataloader_batch_demo.py
   dataset length: 6
   batch features shape: (2, 4)
   batch labels: [5, 3]
   batch features:
   tensor([[20., 21., 22., 23.],
           [12., 13., 14., 15.]])
   epoch 0 label order: [5, 3, 4, 2, 1, 0]
   epoch 1 label order: [5, 3, 4, 2, 1, 0]

3. Check the returned batch shape and target order.

   batch features shape: (2, 4) means the loader returned two samples, each with four feature values. The repeated label order confirms that recreating the seeded Generator makes the shuffled smoke test repeatable.

4. Use the same unpacking pattern in the training loop.

   for batch_features, batch_labels in loader:
       optimizer.zero_grad()
       predictions = model(batch_features)
       loss = loss_fn(predictions, batch_labels)
       loss.backward()
       optimizer.step()

   If the dataset returns one tensor per sample, unpack one value from the batch, such as for (batch_features,) in loader:. Match the loop variables to the dataset fields instead of forcing every dataset into a feature-label pair.

5. Adjust loader options after the single-process batch contract is correct.

   loader = DataLoader(
       dataset,
       batch_size=64,
       shuffle=True,
       num_workers=4,
       pin_memory=torch.cuda.is_available(),
   )

   Keep num_workers=0 while debugging dataset code. Add worker processes only after the loader returns the expected batch shape in the main process.
   Related: How to debug DataLoader worker errors in PyTorch

6. Remove the temporary smoke script.

   $ rm dataloader_batch_demo.py