docstr: Get pipeline layers for current stage.
docstr_tok: ["Get","pipeline","layers","for","current","stage","."]
code: def get_held_layers(self) -> List[Module]:
      """Get pipeline layers for current stage."""
      assert self.pipeline_stage_manager is not None
  
      if self.model.__class__.__name__ == "LlamaModel":
          module = self.model
      else:
          module = self.model.model
      stage_manager = self.pipeline_stage_manager
  
      held_layers = []
      layers_per_stage = self.distribute_layers(
          len(module.layers), stage_manager.num_stages
      )
      if stage_manager.is_first_stage():
          held_layers.append(module.embed_tokens)
          held_layers.append(self.model.lm_head)
      start_idx, end_idx = self.get_stage_index(
          layers_per_stage, stage_manager.stage
      )
      held_layers.extend(module.layers[start_idx:end_idx])
      if stage_manager.is_last_stage():
          held_layers.append(module.norm)
  
      return held_layers
  
file_path: ColossalAI/colossalai/inference/engine/policies/llama.py
imports: from functools import partial
  from typing import List
  import torch
  from torch.nn import Module
  from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaForCausalLM, LlamaModel, LlamaRMSNorm
  from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription
  from colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy
  from ..modeling._utils import init_to_get_rotary
  from ..modeling.llama import LlamaInferenceForwards
label_desc: Use image node_id 4 for calling the LlamaModelInferPolicy obj's underlying member method code with example usage: obj.get_held_layers() and returns: held_layers