docstr: hidden_states: (B, S, D) if batch_first else (S, B, D)
docstr_tok: ["hidden_states",":","(","B",",","S",",","D",")","if","batch_first","else","(","S",",","B",",","D",")"]
code: def forward(self, hidden_states, **kwargs):
      """
      hidden_states: (B, S, D) if batch_first else (S, B, D)
      """
      CausalLMOutput = namedtuple("CausalLMOutput", ["logits"])
      return CausalLMOutput(self.lm_head(hidden_states))
  
imports: import math
  from functools import partial
  from collections import namedtuple
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.nn.modules.utils import _pair
  import hydra
  from einops import reduce, rearrange
label_desc: Use image node_id 2 for calling the LMHead obj's underlying member method code with example usage: obj.forward(hidden_states) and returns: CausalLMOutput