InBatchShuffler

class torchdata.datapipes.iter.InBatchShuffler(datapipe: IterDataPipe[DataChunk[T_co]])

Shuffles each mini-batch from the prior DataPipe (functional name: in_batch_shuffle).

Parameters:

datapipe – Iterable DataPipe with batched data

Example

>>> from torchdata.datapipes.iter import IterableWrapper
>>> source_dp = IterableWrapper(range(10))
>>> batch_dp = source_dp.batch(batch_size=3, drop_last=True)
>>> list(batch_dp)
[[0, 1, 2], [3, 4, 5], [6, 7, 8]]
>>> in_batch_shuffle_dp = batch_dp.in_batch_shuffle()
>>> list(in_batch_shuffle_dp)
[[2, 0, 1], [3, 5, 4], [7, 8, 6]]