class_docstr: Pooler layer.
  
  Pool hidden states of a specific token (for example start of the
  sequence) and add a linear transformation followed by a tanh.
  
  Arguments:
      hidden_size: hidden size
      init_method: weight initialization method for the linear layer.
          bias is set to zero.
class_docstr_tok: ["Pooler","layer",".","Pool","hidden","states","of","a","specific","token","(","for","example","start","of","the","sequence",")","and","add","a","linear","transformation","followed","by","a","tanh",".","Arguments",":","hidden_size",":","hidden","size","init_method",":","weight","initialization","method","for","the","linear","layer",".","bias","is","set","to","zero","."]
code: def forward(self, hidden_states):
      # hidden_states: [b, 1, h]
      # sequence_index: index of the token to pool.
      dense_in_result = self.dense_in(hidden_states)
      tanh_result = torch.tanh(dense_in_result)
      dense_out_result = self.dense_out(tanh_result)
      return dense_out_result
  
file_path: megatron-lm/megatron/model/vision/vit_backbone.py
imports: import math
  import einops
  import torch
  import apex
  import torch.nn.functional
  from megatron import get_args
  from megatron.model.transformer import ParallelTransformer
  from megatron.model.utils import get_linear_layer, init_method_normal, scaled_init_method_normal
  from megatron.model.module import MegatronModule
label_desc: Use image node_id 2 for calling the VitMlpHead obj's underlying member method code with example usage: obj.forward(hidden_states) and returns: dense_out_result