class_docstr: Context manager to manage memory consumption during training.
  
  Allows setting hard limit on the physical batch size as a just one line code change.
  Can be used both for simulating large logical batches with limited memory and for
  safeguarding against occasional large batches produced by
  :class:`~opacus.utils.uniform_sampler.UniformWithReplacementSampler`.
  
  Note that it doesn't modify the input DataLoader, you'd need to use new DataLoader
  returned by the context manager.
  
  BatchSplittingSampler will split large logical batches into smaller sub-batches with
  certain maximum size.
  On every step optimizer will check if the batch was the last physical batch comprising
  a logical one, and will change behaviour accordingly.
  
  If it was not the last, ``optimizer.step()`` will only clip per sample gradients and
  sum them into ``p.summed_grad`.` ``optimizer.zero_grad()`` will clear ``p.grad_sample``,
  but will leave ``p.grad`` and ``p.summed_grad``
  
  If the batch was the last one of the current logical batch, then
  ``optimizer.step()`` and ``optimizer.zero_grad()`` will behave normally.
  
  Example:
      >>> # Assuming you've initialized your objects and passed them to PrivacyEngine.
      >>> # For this example we assume data_loader is initialized with batch_size=4
      >>> model, optimizer, data_loader = _init_private_training()
      >>> criterion = nn.CrossEntropyLoss()
      >>> with BatchMemoryManager(
      ...     data_loader=data_loader, max_physical_batch_size=2, optimizer=optimizer
      ... ) as new_data_loader:
      ...     for data, label in new_data_loader:
      ...         assert len(data) <= 2 # physical batch is no more than 2
      ...         output = model(data)
      ...         loss = criterion(output, label)
      ...         loss.backward()
      ...         # optimizer won't actually make a step unless logical batch is over
      ...         optimizer.step()
      ...         # optimizer won't actually clear gradients unless logical batch is over
      ...         optimizer.zero_grad()
code: def __enter__(self):
      return wrap_data_loader(
          data_loader=self.data_loader,
          max_batch_size=self.max_physical_batch_size,
          optimizer=self.optimizer,
      )
  
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 2 for calling the BatchMemoryManager obj's underlying member method code with example usage: obj.__enter__() and returns: wrap_data_loader