docstr: Load a micro batch from the current batch.
  
  Args:
      microbatch_id (int): the current model chunk idx.
  
  Returns:
      Any: Micro batch.
docstr_tok: ["Load","a","micro","batch","from","the","current","batch",".","Args",":","microbatch_id","(","int",")",":","the","current","model","chunk","idx",".","Returns",":","Any",":","Micro","batch","."]
code: def load_micro_batch(self, model_chunk_id: int) -> Any:
      """Load a micro batch from the current batch.
  
      Args:
          microbatch_id (int): the current model chunk idx.
  
      Returns:
          Any: Micro batch.
      """
      assert (
          self.microbatch_offset[model_chunk_id] <= self.batch_size
      ), "Microbatches exhausted"
      micro_batch = get_micro_batch(
          self.batch,
          self.microbatch_offset[model_chunk_id],
          self.microbatch_size,
      )
      self.microbatch_offset[model_chunk_id] += self.microbatch_size
      return tree_map(
          partial(to_device, device=get_current_device()), micro_batch
      )
  
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 3 for calling the InterleavedSchedule obj's underlying member method code with example usage: obj.load_micro_batch(model_chunk_id) and returns: tree_map