code: def get_held_layers(self) -> List[nn.Module]:
      assert (
          self.pipeline_stage_manager is not None
      ), "pipeline_stage_manager is None"
      stage_manager = self.pipeline_stage_manager
  
      if self.model.__class__.__name__ == "WhisperModel":
          model = self.model
      elif (
          self.model.__class__.__name__
          == "WhisperForConditionalGeneration"
      ):
          model = self.model.model
      else:
          model = None
  
      if model:
          encoder = self.model.get_encoder()
          decoder = self.model.get_decoder()
      else:
          # whisper for audio classification holds encoder only
          encoder = self.model.encoder
          decoder = None
  
      num_encoder_layers = len(encoder.layers)
      if decoder:
          num_decoder_layers = len(decoder.layers)
      else:
          num_decoder_layers = 0
  
      held_layers = []
      (
          layers_per_stage,
          decoder_starting_stage,
      ) = WhisperPolicy.distribute_whisper_layers(
          num_encoder_layers,
          num_decoder_layers,
          stage_manager.num_stages,
      )
      start_idx, end_idx = WhisperPolicy.get_whisper_stage_index(
          layers_per_stage, stage_manager.stage, decoder_starting_stage
      )
  
      if stage_manager.stage < decoder_starting_stage:
          # current stage is in whisper's encoder
          if stage_manager.is_first_stage():
              held_layers.append(encoder.embed_positions)
              held_layers.append(encoder.conv1)
              held_layers.append(encoder.conv2)
          if stage_manager.stage == decoder_starting_stage - 1:
              held_layers.append(encoder.layer_norm)
          held_layers.extend(encoder.layers[start_idx:end_idx])
      else:
          # current stage is in whisper's decoder
          # TODO:(Jianghai) We divide encoder and decoder layers into different parts here,
          # the case encoder and decoder put in same stage should be add in the future.
          if stage_manager.stage == decoder_starting_stage:
              held_layers.append(decoder.embed_tokens)
              held_layers.append(decoder.embed_positions)
          if stage_manager.is_last_stage():
              held_layers.append(decoder.layer_norm)
          held_layers.extend(decoder.layers[start_idx:end_idx])
      return held_layers
  
imports: import warnings
  from functools import partial
  from typing import Callable, Dict, List, Tuple
  import numpy
  import torch.nn
  from torch import Tensor
  import colossalai.shardformer.layer
  from ..modeling.jit import get_jit_fused_dropout_add_func
  from ..modeling.whisper import WhisperPipelineForwards, get_jit_fused_whisper_decoder_layer_forward, get_jit_fused_whisper_encoder_layer_forward, get_whisper_flash_attention_forward
  from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
label_desc: Use image node_id 9 for calling the WhisperPolicy obj's underlying member method code with example usage: obj.get_held_layers() and returns: held_layers