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__ == "BloomModel":
          module = self.model
      else:
          module = self.model.transformer
      stage_manager = self.pipeline_stage_manager
  
      held_layers = []
      layers_per_stage = self.distribute_layers(
          len(module.h), stage_manager.num_stages
      )
      if stage_manager.is_first_stage():
          held_layers.append(module.word_embeddings)
          held_layers.append(module.word_embeddings_layernorm)
          held_layers.append(self.model.lm_head)
      start_idx, end_idx = self.get_stage_index(
          layers_per_stage, stage_manager.stage
      )
      held_layers.extend(module.h[start_idx:end_idx])
      if stage_manager.is_last_stage():
          held_layers.append(module.ln_f)
  
      return held_layers
  
file_path: colossalai/colossalai/inference/engine/policies/bloom.py
imports: from functools import partial
  from typing import List
  import torch
  from torch.nn import LayerNorm, Module
  import colossalai.shardformer.layer
  from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription
  from colossalai.shardformer.policies.bloom import BloomForCausalLMPolicy
  from ..modeling.bloom import BloomInferenceForwards
label_desc: Use image node_id 3 for calling the BloomModelInferPolicy obj's underlying member method code with example usage: obj.get_held_layers() and returns: held_layers