docstr: Get pipeline layers for current stage.
docstr_tok: ["Get","pipeline","layers","for","current","stage","."]
code: def get_held_layers(self) -> List[nn.Module]:
      """Get pipeline layers for current stage."""
      assert self.pipeline_stage_manager is not None
      stage_manager = self.pipeline_stage_manager
  
      model = self.model
      encoder = self.model.encoder
      decoder = getattr(self.model, "decoder", None)
  
      num_encoder_layers = len(encoder.block)
      num_decoder_layers = len(decoder.block) if decoder else 0
  
      held_layers = []
      (
          layers_per_stage,
          decoder_starting_stage,
      ) = T5BasePolicy.distribute_t5_layers(
          num_encoder_layers,
          num_decoder_layers,
          stage_manager.num_stages,
      )
      start_idx, end_idx = T5BasePolicy.get_t5_stage_index(
          layers_per_stage, stage_manager.stage, decoder_starting_stage
      )
  
      if stage_manager.stage < decoder_starting_stage:
          # current stage is in t5's encoder
          if stage_manager.is_first_stage():
              held_layers.append(model.shared)
              held_layers.append(encoder.embed_tokens)
              held_layers.append(encoder.dropout)
          if stage_manager.stage == decoder_starting_stage - 1:
              held_layers.append(encoder.final_layer_norm)
              held_layers.append(encoder.dropout)
          held_layers.extend(encoder.block[start_idx:end_idx])
      else:
          # current stage is in t5's decoder
          if stage_manager.stage == decoder_starting_stage:
              held_layers.append(decoder.embed_tokens)
              held_layers.append(decoder.dropout)
          if stage_manager.is_last_stage():
              held_layers.append(decoder.final_layer_norm)
              held_layers.append(decoder.dropout)
          held_layers.extend(decoder.block[start_idx:end_idx])
      return held_layers
  
imports: import warnings
  from functools import partial
  from typing import Callable, Dict, List, Tuple
  import numpy
  from torch import Tensor, nn
  from colossalai.shardformer.layer import DropoutForParallelInput, Embedding1D, FusedRMSNorm, Linear1D_Col, Linear1D_Row, RMSNorm, VocabParallelEmbedding1D
  from colossalai.shardformer.policies.base_policy import ModulePolicyDescription
  from ..modeling.jit import get_jit_fused_dropout_add_func
  from ..modeling.t5 import T5PipelineForwards, get_jit_fused_T5_layer_ff_forward, get_t5_flash_attention_forward, get_T5_layer_cross_attention_forward, get_T5_layer_self_attention_forward
  from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
label_desc: Use image node_id 7 for calling the T5BasePolicy obj's underlying member method code with example usage: obj.get_held_layers() and returns: held_layers