docstr: Load a batch from data iterator.
  
  Args:
      data_iter (Iterable): Data iterator.
      device (Optional[torch.device], optional): Target device. Defaults to None.
docstr_tok: ["Load","a","batch","from","data","iterator",".","Args",":","data_iter","(","Iterable",")",":","Data","iterator",".","device","(","Optional","[","torch.device","]",",","optional",")",":","Target","device",".","Defaults","to","None","."]
code: def load_batch(
      self, data_iter: Iterable, device: Optional[torch.device] = None
  ) -> None:
      """Load a batch from data iterator.
  
      Args:
          data_iter (Iterable): Data iterator.
          device (Optional[torch.device], optional): Target device. Defaults to None.
      """
      batch = next(data_iter)
      if device is not None:
          batch = tree_map(partial(to_device, device=device), batch)
  
      self.microbatch_offset = [0 for _ in range(self.num_model_chunks)]
      self.batch = batch
      self.batch_size = get_batch_size(batch)
  
      if self.microbatch_size is None:
          assert (
              self.batch_size % self.num_microbatch == 0
          ), "Batch size should divided by the number of microbatch"
          self.microbatch_size = self.batch_size // self.num_microbatch
      if self.num_microbatch is None:
          assert (
              self.batch_size % self.microbatch_size == 0
          ), "Batch size should divided by the microbatch size"
          self.num_microbatch = self.batch_size // self.microbatch_size
  
      if not self.forward_only:
          assert (
              self.last_batch_size is None
              or self.last_batch_size == self.batch_size
          )
          assert (
              self.batch_size
              == self.microbatch_size * self.num_microbatch
          )
  
          assert (
              self.num_microbatch % self.stage_manager.num_stages == 0
          ), "Number of microbatch should be an integer multiple of number of pipeline parallel devices"
  
      if self.forward_only:
          self.num_microbatch = (
              self.batch_size - 1
          ) // self.microbatch_size + 1
          # NOTE: disable metadata cache when batch size changes (not valid anymore)
          if self.batch_size != self.last_batch_size:
              self.enable_metadata_cache = False
              self.send_tensor_metadata = True
              self.send_grad_metadata = True
              self.tensor_metadata_recv = None
              self.grad_metadata_recv = None
  
      self.last_batch_size = self.batch_size
  
imports: from functools import partial
  from typing import Any, Callable, Dict, Iterable, List, Optional, Union
  import torch
  import torch.cuda
  from torch.nn import Module, ModuleList
  from torch.utils._pytree import tree_map
  from colossalai.interface import OptimizerWrapper
  from colossalai.pipeline.p2p import PipelineP2PCommunication, create_send_metadata
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.utils.device import get_current_device
  from ._utils import detach, get_batch_size, get_micro_batch, merge_batch, model_forward, retain_grad, to_device
  from .base import PipelineSchedule
label_desc: Use image node_id 2 for calling the InterleavedSchedule obj's underlying member method code with example usage: obj.load_batch(data_iter, device) without return types