class_docstr: Transformer class.
class_docstr_tok: ["Transformer","class","."]
docstr: Set input tensor to be used instead of forward()'s input.
  
  When doing pipeline parallelism the input from the previous
  stage comes from communication, not from the input, so the
  model's forward_step_func won't have it. This function is thus
  used by internal code to bypass the input provided by the
  forward_step_func
docstr_tok: ["Set","input","tensor","to","be","used","instead","of","forward","(",")","'s","input",".","When","doing","pipeline","parallelism","the","input","from","the","previous","stage","comes","from","communication",",","not","from","the","input",",","so","the","model","'s","forward_step_func","wo","n't","have","it",".","This","function","is","thus","used","by","internal","code","to","bypass","the","input","provided","by","the","forward_step_func"]
code: def set_input_tensor(self, input_tensor: Tensor):
      """Set input tensor to be used instead of forward()'s input.
  
      When doing pipeline parallelism the input from the previous
      stage comes from communication, not from the input, so the
      model's forward_step_func won't have it. This function is thus
      used by internal code to bypass the input provided by the
      forward_step_func"""
      self.input_tensor = input_tensor
  
imports: import re
  from contextlib import nullcontext
  from dataclasses import dataclass
  from typing import List, Union
  import torch
  from torch import Tensor
  from megatron.core import InferenceParams, parallel_state, tensor_parallel
  from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
  from megatron.core.transformer.custom_layers.transformer_engine import TENorm
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.spec_utils import ModuleSpec, build_module
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.transformer.transformer_layer import TransformerLayer
  from megatron.core.utils import make_sharded_tensor_for_checkpoint, make_viewless_tensor
label_desc: Use image node_id 5 for calling the TransformerBlock obj's underlying member method code with example usage: obj.set_input_tensor(input_tensor) without return types