code: def forward(
      self,
      sequence_output: torch.Tensor,
      encoded_layers: Optional[List[torch.Tensor]] = None,
      processed_sample_list: Optional[
          Dict[str, Dict[str, torch.Tensor]]
      ] = None,
  ):
      assert (
          processed_sample_list is not None
      ), "MLM head requires 'processed_sample_list' argument"
  
      output_dict = {}
  
      if self.config.label_key is not None:
          assert self.config.label_key in processed_sample_list, (
              f"Didn't find label key {self.config.label_key} in "
              + "SampleList required by MLM"
          )
          masked_labels = processed_sample_list[self.config.label_key]
      else:
          assert (
              LABEL_KEY in processed_sample_list
              and processed_sample_list[LABEL_KEY] is not None
          ), (
              f"MLM pretraining requires {LABEL_KEY} to be in sample "
              + "list with value not None."
          )
  
          assert (
              COMBINED_LABEL_KEY in processed_sample_list[LABEL_KEY]
          ), f"labels for all modalities must be concatenated in {COMBINED_LABEL_KEY}"
  
          masked_labels = processed_sample_list[LABEL_KEY][
              COMBINED_LABEL_KEY
          ]
  
      masked_tokens = masked_labels.ne(self.config.ignore_index)
  
      masked_labels = masked_labels[masked_tokens]
      sequence_output = sequence_output[masked_tokens, :]
  
      prediction = self.cls(sequence_output)
      output_dict["logits"] = prediction
      masked_lm_loss = self.ce_loss(
          prediction.contiguous().view(-1, self.vocab_size),
          masked_labels.contiguous().view(-1),
      )
      # When masked_labels are all ignore_index then masked_lm_loss is NaN,
      # so we replace NaN with 0.
      if torch.isnan(masked_lm_loss):
          warnings.warn(
              "NaN detected in masked_lm_loss. Replacing it with 0."
          )
          masked_lm_loss = torch.nan_to_num(masked_lm_loss, nan=0.0)
      output_dict["losses"] = {}
      output_dict["losses"][self.config.loss_name] = masked_lm_loss
      return output_dict
  
file_path: mmf/mmf/models/transformers/heads/mlm.py
imports: import warnings
  from dataclasses import dataclass
  from typing import Dict, List, Optional
  import torch
  from mmf.common.registry import registry
  from mmf.models.transformers.base import BaseTransformerHead
label_desc: Use image node_id 3 for calling the MLM obj's underlying member method code with example usage: obj.forward(sequence_output, encoded_layers, processed_sample_list) and returns: output_dict