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."""
      stage_manager = self.pipeline_stage_manager
      held_layers = super().get_held_layers()
      if stage_manager.is_first_stage():
          held_layers.append(self.model.lm_head)
      return held_layers
  
file_path: colossalai/colossalai/legacy/inference/hybridengine/polices/llama.py
imports: from functools import partial
  from typing import List
  import torch
  from torch.nn import Module
  from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaForCausalLM, LlamaModel, LlamaRMSNorm
  from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription
  from colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy
  from ..modeling._utils import init_to_get_rotary
  from ..modeling.llama import LlamaInferenceForwards
label_desc: Use image node_id 4 for calling the LlamaModelInferPolicy obj's underlying member method code with example usage: obj.get_held_layers() and returns: held_layers