code: def get_collate_for_dataset(
      dataset: Union[Dataset, ConcatDataset],
      ensure_collate_fn_are_the_same: bool = True,
  ) -> Callable:
      """
      Return collate_fn function for dataset. By default, default_collate returned.
      If the dataset has method get_collate_fn() we will use it's return value instead.
      If the dataset is ConcatDataset, we will check whether all get_collate_fn() returns
      the same function.
  
      Args:
          dataset: Input dataset
  
      Returns:
          Collate function to put into DataLoader
      """
      collate_fn = default_collate
  
      if hasattr(dataset, "get_collate_fn"):
          return dataset.get_collate_fn()
      elif isinstance(dataset, ConcatDataset):
          collate_fns = [
              get_collate_for_dataset(ds) for ds in dataset.datasets
          ]
          collate_fn = collate_fns[0]
  
          if ensure_collate_fn_are_the_same:
              for other_collate_fn in collate_fns[1:]:
                  if type(other_collate_fn) != type(collate_fn):
                      raise ValueError(
                          f"Detected ConcatDataset consist of datasets with different collate functions: {type(collate_fn)} and {type(other_collate_fn)}."
                      )
  
                  if isinstance(collate_fn, functools.partial):
                      if not _partial_functions_equal(
                          collate_fn, other_collate_fn
                      ):
                          raise ValueError(
                              f"Detected ConcatDataset consist of datasets with different collate functions: {collate_fn} and {type(other_collate_fn)}."
                          )
                  elif collate_fn != other_collate_fn:
                      raise ValueError(
                          f"Detected ConcatDataset consist of datasets with different collate functions: {collate_fn} and {other_collate_fn}."
                      )
  
          collate_fn = collate_fns[0]
  
      return collate_fn
  
imports: import collections
  import dataclasses
  import functools
  import warnings
  import logging
  from typing import Optional, Sequence, Union, Dict, List, Any, Iterable, Callable
  import numpy
  import torch
  from torch import nn, Tensor
  from torch.utils.data import Dataset, ConcatDataset
  from torch.utils.data.dataloader import default_collate
  from .support import pytorch_toolbelt_deprecated
label_desc: Use image node_id 25 for calling a global function with example usage: get_collate_for_dataset(dataset, ensure_collate_fn_are_the_same) and returns: collate_fn, dataset