pytorch
dataloaders
iterate
simultaneous
deep learning

How to iterate over two dataloaders simultaneously using pytorch?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

In many machine learning workflows you need to process two datasets in lockstep. Examples include training a GAN where one loader provides real images and another provides noise vectors, aligning paired datasets for image-to-image translation, or co-training a model on labeled and unlabeled data simultaneously. PyTorch's DataLoader does not natively support iterating over two loaders at once, but Python provides clean patterns for combining them.

This article covers the practical techniques for simultaneous iteration, including handling datasets of different lengths and maintaining proper shuffling behavior.

Quick Recap of DataLoader

A PyTorch DataLoader wraps a Dataset and provides batched, optionally shuffled, and parallelized iteration.

python
1from torch.utils.data import DataLoader, TensorDataset
2import torch
3
4# Two simple datasets
5dataset_a = TensorDataset(torch.randn(1000, 10), torch.randint(0, 2, (1000,)))
6dataset_b = TensorDataset(torch.randn(800, 10), torch.randint(0, 5, (800,)))
7
8loader_a = DataLoader(dataset_a, batch_size=32, shuffle=True)
9loader_b = DataLoader(dataset_b, batch_size=32, shuffle=True)

Notice that dataset_a has 1000 samples and dataset_b has 800. This length mismatch is the main challenge when iterating simultaneously.

Method 1: Using zip (Shortest Dataset Wins)

Python's built-in zip pairs elements from two iterables and stops when the shorter one is exhausted.

python
1for (x_a, y_a), (x_b, y_b) in zip(loader_a, loader_b):
2    # Process both batches
3    combined = torch.cat([x_a, x_b], dim=0)
4    print(f"Batch sizes: {x_a.shape[0]}, {x_b.shape[0]}")
5    # Training step here

This is the simplest approach and works well when both datasets are the same length or when you intentionally want to stop at the shorter one. With the datasets above, iteration stops after 25 batches (800 / 32) even though loader_a has 7 more batches available.

Method 2: Using itertools.zip_longest (All Data Used)

If you want to iterate until the longer dataset is exhausted, use zip_longest and handle the None values from the shorter loader.

python
1from itertools import zip_longest
2
3for batch_a, batch_b in zip_longest(loader_a, loader_b):
4    if batch_a is not None:
5        x_a, y_a = batch_a
6    else:
7        # Handle missing batch from loader_a
8        continue
9
10    if batch_b is not None:
11        x_b, y_b = batch_b
12    else:
13        # Handle missing batch from loader_b
14        continue
15
16    print(f"Batch sizes: {x_a.shape[0]}, {x_b.shape[0]}")

A common strategy for the None case is to skip the batch, use a default tensor, or cycle the shorter loader (see Method 3).

Method 3: Cycling the Shorter Loader

For training loops where you want to reuse the shorter dataset, wrap it with itertools.cycle so it restarts automatically.

python
1from itertools import cycle
2
3# loader_b is shorter, so cycle it
4for (x_a, y_a), (x_b, y_b) in zip(loader_a, cycle(loader_b)):
5    print(f"Batch sizes: {x_a.shape[0]}, {x_b.shape[0]}")

This approach is widely used in semi-supervised learning, where labeled data is scarce and unlabeled data is abundant. The labeled loader cycles through its data multiple times per epoch while the unlabeled loader completes one full pass.

Be aware that cycle does not re-shuffle the shorter dataset when it restarts. If shuffling matters, you need to manually reset the loader (see the custom iterator below).

Method 4: Custom Infinite Loader with Re-Shuffling

To cycle through a shorter dataset with fresh shuffling on each pass, create a custom iterator.

python
1def infinite_loader(dataloader):
2    """Yield batches indefinitely, re-creating the iterator each cycle."""
3    while True:
4        for batch in dataloader:
5            yield batch
6
7loader_b_infinite = infinite_loader(loader_b)
8
9for (x_a, y_a) in loader_a:
10    x_b, y_b = next(loader_b_infinite)
11    print(f"Batch sizes: {x_a.shape[0]}, {x_b.shape[0]}")

Because the for batch in dataloader loop creates a new iterator each time the inner loop ends, the DataLoader re-shuffles the data (assuming shuffle=True was set). This gives you both cycling and proper randomization.

Full Training Loop Example

Here is a complete example that trains a simple model using two loaders for a domain adaptation scenario.

python
1import torch
2import torch.nn as nn
3from torch.utils.data import DataLoader, TensorDataset
4
5# Source domain (labeled)
6src_x = torch.randn(2000, 20)
7src_y = torch.randint(0, 3, (2000,))
8src_loader = DataLoader(TensorDataset(src_x, src_y), batch_size=64, shuffle=True)
9
10# Target domain (unlabeled, using dummy labels)
11tgt_x = torch.randn(1500, 20)
12tgt_y = torch.zeros(1500, dtype=torch.long)  # placeholder
13tgt_loader = DataLoader(TensorDataset(tgt_x, tgt_y), batch_size=64, shuffle=True)
14
15model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 3))
16optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
17criterion = nn.CrossEntropyLoss()
18
19def infinite_loader(dl):
20    while True:
21        for batch in dl:
22            yield batch
23
24tgt_iter = infinite_loader(tgt_loader)
25
26for epoch in range(5):
27    total_loss = 0.0
28    for (sx, sy) in src_loader:
29        tx, _ = next(tgt_iter)
30
31        # Classification loss on source
32        src_pred = model(sx)
33        loss = criterion(src_pred, sy)
34
35        # You could add a domain adaptation loss using tx here
36
37        optimizer.zero_grad()
38        loss.backward()
39        optimizer.step()
40        total_loss += loss.item()
41
42    avg_loss = total_loss / len(src_loader)
43    print(f"Epoch {epoch + 1}, Loss: {avg_loss:.4f}")

Common Pitfalls

Mismatched batch sizes. If the two loaders use different batch sizes, operations that assume equal-sized tensors (like torch.cat along the batch dimension) will still work, but your effective batch size will vary. Make sure downstream code handles variable sizes.

Forgetting that zip silently drops data. When datasets have different lengths, zip stops at the shorter one without any warning. If you do not realize data is being dropped, model performance can degrade. Always log how many batches each loader would produce independently.

Shuffling desynchronization. Both loaders shuffle independently. If your datasets are paired (for example, image A and its corresponding mask), they must come from the same Dataset instance, not two separate loaders. Use a single DataLoader with a dataset that returns both items per sample.

Worker process overhead. Each DataLoader can spawn num_workers processes for parallel loading. Two loaders with num_workers=4 each means 8 worker processes. On memory-constrained machines, reduce worker counts or share a single loader where possible.

Not setting the same random seed. If reproducibility matters and you use shuffle=True on both loaders, set torch.manual_seed() and worker_init_fn consistently to ensure the same pairing across runs.

Summary

To iterate over two PyTorch DataLoaders simultaneously, use zip when both datasets are the same length, zip_longest when you need all data from both, or cycle/infinite_loader when the shorter dataset should repeat. For paired data, always use a single DataLoader with a combined Dataset rather than two separate loaders. The infinite_loader pattern with a generator function is the most flexible approach because it supports re-shuffling on each cycle.


Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.