class_docstr: GenerateSchedule is a class that handles the pipeline parallel inference.
  In our schedule, we place tie weight layer, embedding and lm_head in the same device to save space, so in
  this schedule, the out for each encoding progress is on rank0.
  
  Args:
      stage_manager (`PipelineStageManager`): Pipeline stage manager.
      mb_manager (`MicroBatchManager`): Micro batch manager.
      verbose (bool): Whether to verbose the information of the pipeline.
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.batch = batch
      self.batch_size = get_batch_size(batch)
      if self.stage_manager.num_stages == 1:
          self.microbatch_size = self.batch_size
      self.microbatch_offset = 0
      assert (
          self.batch_size % self.microbatch_size == 0
      ), f"Batch size should divided by the number of microbatches, {self.batch_size}, {self.num_microbatches}"
      self.num_microbatches = self.batch_size // self.microbatch_size
      self.round = (
          self.num_microbatches // self.stage_manager.num_stages
      )
  
imports: import time
  from functools import partial
  from typing import Any, Iterable, Optional, Union
  import torch
  import torch.cuda
  from torch.nn import Module
  from torch.utils._pytree import tree_map
  from colossalai.inference.engine.microbatch_manager import MicroBatchManager, Status
  from colossalai.pipeline.p2p import PipelineP2PCommunication
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.utils.device import get_current_device
  from ._utils import get_batch_size, get_micro_batch, model_forward, to_device
  from .base import PipelineSchedule
label_desc: Use image node_id 2 for calling the GenerateSchedule obj's underlying member method code with example usage: obj.load_batch(data_iter, device) without return types