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."""
      assert self.pipeline_stage_manager is not None
  
      if self.model.__class__.__name__ == "BertModel":
          module = self.model
      else:
          module = self.model.bert
      stage_manager = self.pipeline_stage_manager
  
      held_layers = []
      if stage_manager.is_interleave:
          assert stage_manager.num_model_chunks is not None
          layers_per_stage = self.distribute_layers(
              len(module.encoder.layer),
              stage_manager.num_stages * stage_manager.num_model_chunks,
          )
          stage_indices = Policy.get_stage_index(
              layers_per_stage,
              stage_manager.stage,
              num_model_chunks=stage_manager.num_model_chunks,
              num_stages=stage_manager.num_stages,
          )
          if stage_manager.is_first_stage(ignore_chunk=True):
              held_layers.append(module.embeddings)
          for start_idx, end_idx in stage_indices:
              held_layers.extend(
                  module.encoder.layer[start_idx:end_idx]
              )
          if stage_manager.is_last_stage(ignore_chunk=True):
              held_layers.append(module.pooler)
  
      else:
          layers_per_stage = self.distribute_layers(
              len(module.encoder.layer), stage_manager.num_stages
          )
          if stage_manager.is_first_stage():
              held_layers.append(module.embeddings)
          start_idx, end_idx = Policy.get_stage_index(
              layers_per_stage, stage_manager.stage
          )
          held_layers.extend(module.encoder.layer[start_idx:end_idx])
          if stage_manager.is_last_stage():
              held_layers.append(module.pooler)
  
      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 8 for calling the BertPolicy obj's underlying member method code with example usage: obj.get_held_layers() and returns: held_layers