class_docstr: A basic DataModule class for all Flash tasks. This class includes references to a
  :class:`~flash.core.data.datasets.Input` and a :class:`~flash.core.data.callback.BaseDataFetcher`.
  
  Args:
      train_input: Input dataset for training. Defaults to None.
      val_input: Input dataset for validating model performance during training. Defaults to None.
      test_input: Input dataset to test model performance. Defaults to None.
      predict_input: Input dataset for predicting. Defaults to None.
      data_fetcher: The :class:`~flash.core.data.callback.BaseDataFetcher` to attach to the
          :class:`~flash.core.data.io.input_transform.InputTransform`. If ``None``, the output from
          :meth:`~flash.core.data.data_module.DataModule.configure_data_fetcher` will be used.
      transform: The :class:`~flash.core.data.io.input_transform.InputTransform` type to use.
      transform_kwargs: Dict of keyword arguments to be provided when instantiating the transforms.
      val_split: An optional float which gives the relative amount of the training dataset to use for the validation
          dataset.
      batch_size: The batch size to be used by the DataLoader.
      num_workers: The number of workers to use for parallelized loading.
      sampler: A sampler following the :class:`~torch.utils.data.sampler.Sampler` type.
          Will be passed to the DataLoader for the training dataset. Defaults to None.
  
  Examples
  ________
  
  .. testsetup::
  
      >>> from flash import DataModule
      >>> from flash.core.utilities.stages import RunningStage
      >>> from torch.utils.data.sampler import SequentialSampler, WeightedRandomSampler
      >>> class TestInput(Input):
      ...     def train_load_data(self, _):
      ...         return [(0, 1, 2, 3), (0, 1, 2, 3)]
      >>> train_input = TestInput(RunningStage.TRAINING, [1])
  
  You can provide the sampler to use for the train dataloader using the ``sampler`` argument.
  The sampler can be a function or type that needs the dataset as an argument:
  
  .. doctest::
  
      >>> datamodule = DataModule(train_input, sampler=SequentialSampler, batch_size=1)
      >>> print(datamodule.train_dataloader().sampler)  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
      <torch.utils.data.sampler.SequentialSampler object at ...>
  
  Alternatively, you can pass a sampler instance:
  
  .. doctest::
  
      >>> datamodule = DataModule(train_input, sampler=WeightedRandomSampler([0.1, 0.5], 2), batch_size=1)
      >>> print(datamodule.train_dataloader().sampler)  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
      <torch.utils.data.sampler.WeightedRandomSampler object at ...>
code: def __init__(
      self,
      train_input: Optional[Input] = None,
      val_input: Optional[Input] = None,
      test_input: Optional[Input] = None,
      predict_input: Optional[Input] = None,
      data_fetcher: Optional[BaseDataFetcher] = None,
      transform: INPUT_TRANSFORM_TYPE = InputTransform,
      transform_kwargs: Optional[Dict] = None,
      val_split: Optional[float] = None,
      batch_size: Optional[int] = None,
      num_workers: int = 0,
      sampler: Optional[Union[Callable, Sampler, Type[Sampler]]] = None,
      pin_memory: bool = True,
      persistent_workers: bool = False,
  ) -> None:
      if not batch_size:
          raise TypeError(
              "The `batch_size` should be provided to the DataModule on instantiation."
          )
  
      if flash._IS_TESTING and torch.cuda.is_available():
          batch_size = 16
  
      self.input_transform = create_or_configure_input_transform(
          transform=transform, transform_kwargs=transform_kwargs
      )
  
      self.viz: Optional[BaseVisualization] = None
  
      self._train_input = train_input
      self._val_input = val_input
      self._test_input = test_input
      self._predict_input = predict_input
  
      if (
          self._train_input
          and self._val_input
          and isinstance(val_split, float)
          and val_split > 0
      ):
          raise TypeError(
              "A `val_dataset` was provided with `val_split`. Please, choose one or the other."
          )
  
      if self._train_input and (
          val_split is not None and not self._val_input
      ):
          self._train_input, self._val_input = self._split_train_val(
              self._train_input, val_split
          )
  
      self._data_fetcher: Optional[BaseDataFetcher] = (
          data_fetcher or self.configure_data_fetcher()
      )
  
      self._on_after_batch_transfer_fns = None
  
      if self._train_input:
          self.train_dataloader = self._train_dataloader
  
      if self._val_input:
          self.val_dataloader = self._val_dataloader
  
      if self._test_input:
          self.test_dataloader = self._test_dataloader
  
      if self._predict_input:
          self.predict_dataloader = self._predict_dataloader
  
      self.batch_size = batch_size
  
      self.num_workers = num_workers
      self.persistent_workers = persistent_workers and num_workers > 0
      self.pin_memory = pin_memory
  
      self.sampler = sampler
  
      super().__init__()
  
imports: import os
  from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type, Union
  import numpy
  import pytorch_lightning
  import torch
  from torch.utils.data import DataLoader, Dataset
  from torch.utils.data.dataset import IterableDataset
  from torch.utils.data.sampler import Sampler
  import flash
  from flash.core.data.base_viz import BaseVisualization
  from flash.core.data.callback import BaseDataFetcher
  from flash.core.data.io.input import DataKeys, Input, IterableInput
  from flash.core.data.io.input_transform import InputTransform, create_device_input_transform_processor, create_or_configure_input_transform, create_worker_input_transform_processor
  from flash.core.data.splits import SplitDataset
  from flash.core.data.utils import _STAGES_PREFIX
  from flash.core.utilities.imports import _TOPIC_CORE_AVAILABLE
  from flash.core.utilities.stages import RunningStage
  from flash.core.utilities.types import INPUT_TRANSFORM_TYPE
label_desc: Use image node_id 1 to create a new DataModule object from inherited base classes: pl with example: obj = DataModule(train_input, val_input, test_input, predict_input, data_fetcher, transform, transform_kwargs, val_split, batch_size, num_workers, sampler, pin_memory, persistent_workers)