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.
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.
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.
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.
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.
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.
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.
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
- How to iterate through tensors in custom loss function?
- How to let TensorFlow XLA know the CUDA path
- how to limit GPU usage in tensorflow r1.1 with C API
- How to load a model from an HDF5 file in Keras?
- How to parallelize a training loop ever samples of a batch when CPU is only available in pytorch?
- How to reverse the operation of torch.nn.functional.grid_sample?
- How to load a tflite model in script?
- How to load a trained model''s weights, which were saved with tf.keras.models.save_model?
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.