code: def get_sampler(dataset, dataset_config, sampler_seed=0):
      """
      Given the dataset object and the dataset config, get the data sampler to use
      Supports 2 types of samplers:
          - Pytorch default torch.utils.data.distributed.DistributedSampler
          - VISSL sampler StatefulDistributedSampler that is written specifically for
            large scale dataset trainings
      """
      data_sampler = None
      if (
          torch.distributed.is_available()
          and torch.distributed.is_initialized()
      ):
          if dataset_config["STRATIFIED_SAMPLER"]["ENABLED"]:
              class_index, labels = create_class_index(dataset)
              min_samples_per_class = min(
                  len(v) for k, v in class_index.items()
              )
              local_batch_size = dataset_config["BATCHSIZE_PER_REPLICA"]
              step_per_epoch = dataset_config["STRATIFIED_SAMPLER"][
                  "STEP_PER_EPOCH"
              ]
              classes_per_batch = dataset_config["STRATIFIED_SAMPLER"][
                  "CLASSES_PER_BATCH"
              ]
              unique_classes = dataset_config["STRATIFIED_SAMPLER"][
                  "UNIQUE_CLASSES"
              ]
              world_size = dist.get_world_size()
  
              # In case each worker generates the same classes (not unique worker)
              if not unique_classes:
                  samples_per_batch, r = divmod(
                      local_batch_size, classes_per_batch
                  )
                  assert (
                      r == 0
                  ), "Number of classes per batch has to divide the local batch size"
  
                  global_samples_per_batch = (
                      world_size * samples_per_batch
                  )
                  error_msg = f"Not enough classes per batch to build a global batch: {global_samples_per_batch}"
                  assert (
                      global_samples_per_batch < min_samples_per_class
                  ), error_msg
  
                  data_sampler = StratifiedClassSampler(
                      dataset,
                      num_classes=len(class_index),
                      class_index=class_index,
                      labels=labels,
                      world_size=world_size,
                      rank=dist.get_rank(),
                      batch_size=samples_per_batch,
                      classes_per_batch=classes_per_batch,
                      epochs=1,
                      seed=0,
                      unique_classes=unique_classes,
                  )
  
              # In case each worker generates different uniques classes
              else:
                  classes_per_worker, r = divmod(
                      classes_per_batch, world_size
                  )
                  assert (
                      r == 0
                  ), f"Number of classes per batch {classes_per_batch} has to divide the world size {world_size}"
  
                  sampler_batch_size, r = divmod(
                      local_batch_size, classes_per_worker
                  )
                  assert (
                      r == 0
                  ), f"Number of classes per local batch {classes_per_worker} has to divide the local batch size {local_batch_size}"
  
                  error_msg = f"Not enough classes per batch to build a global batch: {sampler_batch_size}"
                  assert (
                      sampler_batch_size < min_samples_per_class
                  ), error_msg
  
                  data_sampler = StratifiedClassSampler(
                      dataset,
                      num_classes=len(class_index),
                      class_index=class_index,
                      labels=labels,
                      world_size=world_size,
                      rank=dist.get_rank(),
                      batch_size=sampler_batch_size,
                      classes_per_batch=classes_per_worker,
                      epochs=1,
                      seed=0,
                      unique_classes=unique_classes,
                  )
  
              data_sampler.set_inner_epochs(
                  step_per_epoch // len(data_sampler)
              )
          elif dataset_config["POWER_LAW_SAMPLER"]["ENABLED"]:
              class_index, labels = create_class_index(dataset)
              local_batch_size = dataset_config["BATCHSIZE_PER_REPLICA"]
              data_sampler = ClassPowerlawSampler(
                  dataset=dataset,
                  num_classes=len(class_index),
                  class_index=class_index,
                  labels=labels,
                  world_size=dist.get_world_size(),
                  rank=dist.get_rank(),
                  batch_size=local_batch_size,
                  powerlaw=dataset_config["POWER_LAW_SAMPLER"]["POWER"],
                  seed=0,
              )
          elif dataset_config["USE_DEBUGGING_SAMPLER"]:
              data_sampler = DeterministicDistributedSampler(
                  dataset,
                  batch_size=dataset_config["BATCHSIZE_PER_REPLICA"],
              )
          elif dataset_config["USE_STATEFUL_DISTRIBUTED_SAMPLER"]:
              data_sampler = StatefulDistributedSampler(
                  dataset,
                  batch_size=dataset_config["BATCHSIZE_PER_REPLICA"],
                  seed=sampler_seed,
              )
          else:
              data_sampler = (
                  torch.utils.data.distributed.DistributedSampler(
                      dataset
                  )
              )
          logging.info("Created the Distributed Sampler....")
          print_sampler_config(data_sampler)
      else:
          logging.warning(
              "Distributed trainer not initialized. Not using the sampler and data will NOT be shuffled"  # NOQA
          )
      return data_sampler
  
imports: import logging
  import random
  import numpy
  import torch
  import torch.distributed
  from classy_vision.dataset import DataloaderAsyncGPUWrapper
  from torch.utils.data import DataLoader
  from vissl.data.airstore_dataset import AirstoreDataset
  from vissl.data.collators import get_collator
  from vissl.data.data_helper import ClassPowerlawSampler, DeterministicDistributedSampler, StatefulDistributedSampler, StratifiedClassSampler
  from vissl.data.dataloader_sync_gpu_wrapper import DataloaderSyncGPUWrapper
  from vissl.data.dataset_catalog import get_data_files, register_datasets, VisslDatasetCatalog
  from vissl.data.disk_dataset import DiskImageDataset
  from vissl.data.ssl_dataset import GenericSSLDataset
  from vissl.data.synthetic_dataset import SyntheticImageDataset
  from vissl.data.torchvision_dataset import TorchvisionDataset
  from vissl.utils.misc import set_dataloader_seeds, setup_multiprocessing_method
label_desc: Use image node_id 5 for calling a global function with example usage: get_sampler(dataset, dataset_config, sampler_seed) and returns: data_sampler