def _show_batch(
        self,
        stage: str,
        func_names: Union[str, List[str]],
        limit_nb_samples: int = None,
        figsize: Tuple[int, int] = (6.4, 4.8),
        reset: bool = True,
    ) -> None:
        """This function is used to handle transforms profiling for batch visualization."""
        # don't show in CI
        if os.getenv("FLASH_TESTING", "0") == "1":
            return
        iter_name = f"_{stage}_iter"

        if not hasattr(self, iter_name):
            self._reset_iterator(stage)

        # list of functions to visualise
        if isinstance(func_names, str):
            func_names = [func_names]

        if not limit_nb_samples:
            limit_nb_samples = self.batch_size

        iter_dataloader = getattr(self, iter_name)
        with self.data_fetcher.enable():
            if reset:
                self.data_fetcher.batches[stage] = {}
            try:
                _ = next(iter_dataloader)
            except StopIteration:
                iter_dataloader = self._reset_iterator(stage)
                _ = next(iter_dataloader)
            data_fetcher: BaseVisualization = self.data_fetcher
            data_fetcher._show(stage, func_names, limit_nb_samples, figsize)
            if reset:
                self.data_fetcher.batches[stage] = {}
This function is used to handle transforms profiling for batch visualization.
src/flash/core/data/data_module.py