How to load/fetch the next data batches for the next epoch, during the current epoch?
I know that since PyTorch 1.7.0 it is possible to prefetch some batches before an epoch begins. However, this does not make it possible to fetch batches while the operations within an epoch are being performed and before the next epoch begins. Based on this thread, it seems that it should be possible to use a Sampler to load batches during an epoch, and before the next epoch begins. However, I cannot wrap my head around how I can use a Sampler to achieve this.
Can anyone provide a code sample for a Sampler that allows fetching samples during an epoch?

You can prefetch the next batches from iterator in a background thread.
class _ThreadedIterator(threading.Thread):
    &quot;&quot;&quot;
    Prefetch the next queue_length items from iterator in a background thread.

    Example:
    &gt;&gt; for i in bg_iterator(range(10)):
    &gt;&gt;     print(i)
    &quot;&quot;&quot;

    class _End:
        pass

    def __init__(self, generator: Iterable, maxsize: int) -&gt; None:
        threading.Thread.__init__(self)
        self.queue: Queue = Queue(maxsize)
        self.generator = generator
        self.daemon = True
        self.start()

    def run(self) -&gt; None:
        for item in self.generator:
            self.queue.put(item)
        self.queue.put(self._End)

    def __iter__(self) -&gt; Any:
        return self

    def __next__(self) -&gt; Any:
        next_item = self.queue.get()
        if next_item == self._End:
            raise StopIteration
        return next_item

    # Required for Python 2.7 compatibility
    def next(self) -&gt; Any:
        return self.__next__()


def bg_iterator(iterable: Iterable, maxsize: int) -&gt; Any:
    return _ThreadedIterator(iterable, maxsize=maxsize)

UPD.
Usage:
model = model.to(device, non_blocking=True)
for inputs, targets in bg_iterator(data_loader, maxsize=2):
    inputs = inputs.to(device, non_blocking=True)
    targets = targets.to(device, non_blocking=True)

example
