def METHOD_NAME(self, epoch: int, shuffle: bool = None) -> DataLoader:
    if shuffle is None:
        shuffle = self.shuffle
    if self.num_iters_per_epoch is not None:
        N = len(self.sampler)
        # If corpus size is larger than the num_per_epoch
        if self.num_iters_per_epoch < N:
            N = len(self.sampler)
            real_epoch, offset = divmod(self.num_iters_per_epoch * epoch, N)
            if offset >= self.num_iters_per_epoch:
                current_batches = self.sampler.generate(real_epoch + self.seed)
                if shuffle:
                    np.random.RandomState(real_epoch + self.seed).shuffle(
                        current_batches
                    )
                batches = current_batches[
                    offset - self.num_iters_per_epoch : offset
                ]
            else:
                prev_batches = self.sampler.generate(real_epoch - 1 + self.seed)
                current_batches = self.sampler.generate(real_epoch + self.seed)
                if shuffle:
                    np.random.RandomState(real_epoch - 1 + self.seed).shuffle(
                        prev_batches
                    )
                    np.random.RandomState(real_epoch + self.seed).shuffle(
                        current_batches
                    )
                batches = (
                    prev_batches[offset - self.num_iters_per_epoch :]
                    + current_batches[:offset]
                )
        # If corpus size is less than the num_per_epoch
        else:
            _epoch, _cursor = divmod(self.num_iters_per_epoch * (epoch - 1), N)
            _remain = self.num_iters_per_epoch
            batches = []
            current_batches = self.sampler.generate(_epoch + self.seed)
            if shuffle:
                np.random.RandomState(_epoch + self.seed).shuffle(current_batches)
            while _remain > 0:
                _batches = current_batches[_cursor : _cursor + _remain]
                batches += _batches
                if _cursor + _remain >= N:
                    _epoch += 1
                    _cursor = 0
                    current_batches = self.sampler.generate(_epoch + self.seed)
                    if shuffle:
                        np.random.RandomState(_epoch + self.seed).shuffle(
                            current_batches
                        )
                else:
                    _cursor = _cursor + _remain
                _remain -= len(_batches)
            assert len(batches) == self.num_iters_per_epoch
    else:
        batches = self.sampler.generate(epoch + self.seed)
        if shuffle:
            np.random.RandomState(epoch + self.seed).shuffle(batches)
    # For backward compatibility for pytorch DataLoader
    if self.collate_fn is not None:
        kwargs = dict(collate_fn=self.collate_fn)
    else:
        kwargs = {}
    # reshuffle whole 'batches' so that elements within a batch can move
    # between different batches
    if self.shuffle_within_batch:
        _bs = len(batches[0])
        batches = list(itertools.chain(*batches))
        np.random.RandomState(epoch + self.seed).shuffle(batches)
        _batches = []
        for ii in range(0, len(batches), _bs):
            _batches.append(batches[ii : ii + _bs])
        batches = _batches
        del _batches
    return DataLoader(
        dataset=self.dataset,
        batch_sampler=batches,
        num_workers=self.num_workers,
        pin_memory=self.pin_memory,
        worker_init_fn=partial(worker_init_fn, base_seed=epoch + self.seed),
        **kwargs,
    )