code: def post_language_model_processing(
      lm_output,
      pooled_output,
      lm_head,
      binary_head,
      lm_labels,
      logit_weights,
      fp16_lm_cross_entropy,
  ):
      # lm_logits: [s, b, vocab_size]
      lm_logits = lm_head(lm_output, logit_weights)
  
      binary_logits = None
      if binary_head is not None:
          # binary_logits: [s, b, 2] or [s, b, vocab_size] if binary_head is Identity
          binary_logits = binary_head(pooled_output)
  
      if lm_labels is None:
          return lm_logits, binary_logits
      else:
          # match shape of labels to lm_logits
          # lm_labels: [b, s] -> [s, b]
          lm_labels = lm_labels.transpose(0, 1).contiguous()
          if fp16_lm_cross_entropy:
              assert lm_logits.dtype == torch.half
              lm_loss = tensor_parallel.vocab_parallel_cross_entropy(
                  lm_logits, lm_labels
              )
          else:
              lm_loss = tensor_parallel.vocab_parallel_cross_entropy(
                  lm_logits.float(), lm_labels
              )
          # lm_loss: [s, b]
          return lm_loss, binary_logits
  
file_path: NeMo/nemo/collections/nlp/models/language_modeling/megatron/bert_model.py
imports: import torch
  from nemo.collections.nlp.modules.common.megatron.language_model import get_language_model
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.transformer import get_layer_norm
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, build_position_ids, erf_gelu, get_linear_layer, init_method_normal, openai_gelu, parallel_lm_logits, scaled_init_method_normal
label_desc: Use image node_id 2 for calling a global function with example usage: post_language_model_processing(lm_output, pooled_output, lm_head, binary_head, lm_labels, logit_weights, fp16_lm_cross_entropy) and returns: lm_logits,  binary_logits, lm_loss,  binary_logits