Diff Coverage

        try:
            start_epoch = self.state.epoch
            for epoch in range(start_epoch, self.train_epochs):
                if epoch != start_epoch:
                    self.train_dataloader.set_epoch(epoch)
                    self.data_iterator = HyperIter(
                        self.train_dataloader, use_background_prefetcher=config.dataloader.use_background_prefetcher
                    )
                self.state.epoch = epoch

                for _ in range(start_step, train_steps):
                    try:
                        self.train_step(self.data_iterator)
                    except StopIteration:
                        logger.info(
                            "epoch:%s Dataloader finished with drop_last %s",
                            epoch,
                            config.dataloader.drop_last,
                        )
                        break

                self.on_epoch_end()
                self.state.epoch = epoch + 1

                print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")

                if config.dataloader.use_background_prefetcher:
                    self.data_iterator.stop()

            self.on_train_end()
        finally:
            if config.dataloader.use_background_prefetcher:
                self.data_iterator.stop()

                # don't mutate the captured state in-place. The underlying dataloader's
                # state_dict() should handle deepcopying if necessary.
                state = self.original_state_dict() if self.original_state_dict else None
                if not self._put_result((item, state)):
                    break
        # The worker must transfer any producer failure back to the training thread.
        except Exception as exc:  # pylint: disable=broad-exception-caught
            self._put_result((exc, None))
        finally:
            # A timed-out stop cannot cancel next(), so the worker must release
            # its own references when that call eventually returns.
            if self.stop_event.is_set():

        while not self.stop_event.is_set():
            try:
                self.queue.put(result, timeout=0.1)
                return True
            except queue.Full:
                continue
        return False

    def _drain_queue(self) -> None:
        """Release every result currently retained by the prefetch queue."""
        while True: