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
      """
      held_layers = super().get_held_layers()
      stage_manager = self.pipeline_stage_manager
      if stage_manager.is_last_stage(ignore_chunk=True):
          held_layers.append(self.model.dropout)
          held_layers.append(self.model.classifier)
      return held_layers
  
imports: from functools import partial
  from typing import Callable, Dict, List
  import torch.nn
  from torch import Tensor
  from torch.nn import Module
  import colossalai.shardformer.layer
  from ..modeling.bert import BertPipelineForwards, bert_sequence_parallel_forward_fn, get_bert_flash_attention_forward, get_jit_fused_bert_output_forward, get_jit_fused_bert_self_output_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 BertForTokenClassificationPolicy obj's underlying member method code with example usage: obj.get_held_layers() and returns: held_layers