code: def get_held_layers(self) -> List[nn.Module]:
      held_layers = super().get_held_layers()
      if self.pipeline_stage_manager.is_last_stage():
          held_layers.append(self.model.transformer.output_layer)
      return held_layers
  
imports: from functools import partial
  from typing import Callable, Dict, List, Union
  import torch.nn
  from torch import Tensor
  import colossalai.shardformer.layer
  from colossalai.shardformer.modeling.chatglm2 import ChatGLMPipelineForwards
  from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel
  from ..modeling.chatglm2 import get_chatglm_sequence_parallel_forward_fn, get_flash_core_attention_forward, get_jit_fused_glm_block_forward
  from ..modeling.jit import get_jit_fused_dropout_add_func
  from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
label_desc: Use image node_id 2 for calling the ChatGLMForConditionalGenerationPolicy obj's underlying member method code with example usage: obj.get_held_layers() and returns: held_layers