code: def test_uniter_for_pretraining(self):
      # UNITER pretraining has 5 pretraining tasks,
      # we have one unique head for each, and in each
      # forward pass we train on a different task.
      # In this test we try running a forward pass
      # through each head.
      heads = {
          "mlm": {"type": "mlm"},
          "itm": {"type": "itm"},
          "mrc": {"type": "mrc"},
          "mrfr": {"type": "mrfr"},
          "wra": {"type": "wra"},
      }
      tasks = "mlm,itm,mrc,mrfr,wra"
      mask_probability = 0.15
      model = UNITERForPretraining(
          head_configs=heads,
          tasks=tasks,
          mask_probability=mask_probability,
      )
      model.eval()
      model = model.to(get_current_device())
      sample_list = self._get_sample_list()
      self._enhance_sample_list_for_pretraining(sample_list)
  
      expected_loss_names = {
          "mlm": "masked_lm_loss",
          "itm": "itm_loss",
          "mrc": "mrc_loss",
          "mrfr": "mrfr_loss",
          "wra": "wra_loss",
      }
  
      for task_name, loss_name in expected_loss_names.items():
          sample_list["task"] = task_name
          with torch.no_grad():
              model_output = model(sample_list)
  
          self.assertTrue("losses" in model_output)
          self.assertTrue(loss_name in model_output["losses"])
  
imports: import gc
  import unittest
  import tests.test_utils
  import torch
  from mmf.common.sample import SampleList
  from mmf.models.uniter import UNITERForClassification, UNITERForPretraining, UNITERImageEmbeddings, UNITERModelBase
  from mmf.utils.build import build_model
  from mmf.utils.configuration import Configuration
  from mmf.utils.env import setup_imports, teardown_imports
  from mmf.utils.general import get_current_device
  from omegaconf import OmegaConf
label_desc: Use image node_id 4 for calling the TestUniterWithHeads obj's underlying member method code with example usage: obj.test_uniter_for_pretraining() without return types