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 ...>
docstr: This property returns the test dataset.
docstr_tok: ["This","property","returns","the","test","dataset","."]
code: def test_dataset(self) -> Optional[Input]:
      """This property returns the test dataset."""
      return self._test_input
  
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 4 for calling the DataModule obj's underlying member method code with example usage: obj.test_dataset() and returns: self