docstr: Helper method to get the model chunk ID given the iteration number.
  
  Args:
      microbatch_id (int): the current microbatch idx
      forward (bool): if is the forward process
  
  Returns:
      int: The model chunk idx of the input microbatch_id
docstr_tok: ["Helper","method","to","get","the","model","chunk","ID","given","the","iteration","number",".","Args",":","microbatch_id","(","int",")",":","the","current","microbatch","idx","forward","(","bool",")",":","if","is","the","forward","process","Returns",":","int",":","The","model","chunk","idx","of","the","input","microbatch_id"]
code: def get_model_chunk_id(
      self, microbatch_id: int, is_forward: bool
  ) -> int:
      """Helper method to get the model chunk ID given the iteration number.
  
      Args:
          microbatch_id (int): the current microbatch idx
          forward (bool): if is the forward process
  
      Returns:
          int: The model chunk idx of the input microbatch_id
      """
      assert microbatch_id < self.num_microbatch * self.num_model_chunks
      microbatch_id_in_group = microbatch_id % (
          self.stage_manager.num_stages * self.num_model_chunks
      )
      model_chunk_id = (
          microbatch_id_in_group // self.stage_manager.num_stages
      )
      if not is_forward:
          model_chunk_id = self.num_model_chunks - model_chunk_id - 1
      return model_chunk_id
  
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 4 for calling the InterleavedSchedule obj's underlying member method code with example usage: obj.get_model_chunk_id(microbatch_id, is_forward) and returns: model_chunk_id