code: def get_held_layers(self) -> List[nn.Module]:
      module = self.model.transformer
      stage_manager = self.pipeline_stage_manager
  
      held_layers = []
      layers_per_stage = self.distribute_layers(
          module.num_layers, stage_manager.num_stages
      )
      if stage_manager.is_first_stage():
          held_layers.append(module.embedding)
          held_layers.append(module.output_layer)
      start_idx, end_idx = self.get_stage_index(
          layers_per_stage, stage_manager.stage
      )
      held_layers.extend(module.encoder.layers[start_idx:end_idx])
      if stage_manager.is_last_stage():
          if module.encoder.post_layer_norm:
              held_layers.append(module.encoder.final_layernorm)
  
      # rotary_pos_emb is needed for all stages
      held_layers.append(module.rotary_pos_emb)
  
      return held_layers
  
file_path: colossalai/colossalai/inference/engine/policies/chatglm2.py
imports: from typing import List
  import torch.nn
  from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel, GLMBlock, GLMTransformer, SelfAttention
  from colossalai.shardformer.policies.chatglm2 import ChatGLMModelPolicy
  from ..modeling._utils import init_to_get_rotary
  from ..modeling.chatglm2 import ChatGLM2InferenceForwards
label_desc: Use image node_id 3 for calling the ChatGLM2InferPolicy obj's underlying member method code with example usage: obj.get_held_layers() and returns: held_layers