def test_loss_fn(
      self,
      separate_memory_values=False,
      num_intermediate_layers=None,
      el_second_im_weight=0.0,
  ):
    """Test loss function runs and produces expected values."""
    config = copy.deepcopy(self.config)
    config['model_config']['encoder_config'][
        'separate_memory_values'] = separate_memory_values
    config['model_config']['encoder_config'][
        'num_intermediate_layers'] = num_intermediate_layers
    config['el_second_im_weight'] = el_second_im_weight
    self.run_model(config, entity_vocab_size=1000)
Test loss function runs and produces expected values.
language/mentionmemory/tasks/mention_memory_task_test.py