code: def wrap_data_loader(
      *,
      data_loader: DataLoader,
      max_batch_size: int,
      optimizer: DPOptimizer,
  ):
      """
      Replaces batch_sampler in the input data loader with ``BatchSplittingSampler``
  
      Args:
          data_loader: Wrapper DataLoader
          max_batch_size: max physical batch size we want to emit
          optimizer: DPOptimizer instance used for training
  
      Returns:
          New DataLoader instance with batch_sampler wrapped in ``BatchSplittingSampler``
      """
  
      return DataLoader(
          dataset=data_loader.dataset,
          batch_sampler=BatchSplittingSampler(
              sampler=data_loader.batch_sampler,
              max_batch_size=max_batch_size,
              optimizer=optimizer,
          ),
          num_workers=data_loader.num_workers,
          collate_fn=data_loader.collate_fn,
          pin_memory=data_loader.pin_memory,
          timeout=data_loader.timeout,
          worker_init_fn=data_loader.worker_init_fn,
          multiprocessing_context=data_loader.multiprocessing_context,
          generator=data_loader.generator,
          prefetch_factor=data_loader.prefetch_factor,
          persistent_workers=data_loader.persistent_workers,
      )
  
file_path: opacus/opacus/utils/batch_memory_manager.py
imports: import math
  from typing import List
  import numpy
  from opacus.optimizers import DPOptimizer
  from opacus.utils.uniform_sampler import DistributedUniformWithReplacementSampler, UniformWithReplacementSampler
  from torch.utils.data import BatchSampler, DataLoader, Sampler
label_desc: Use image node_id 1 for calling a global function with example usage: wrap_data_loader() and returns: DataLoader