def test_loss_fn(
      self,
      separate_memory_values,
      qa_k_top_device,
      qa_k_top_post_selection,
  ):
    """Test loss function runs and produces expected values."""
    config = copy.deepcopy(self.config)
    config['model_config']['encoder_config']['k_top_device'] = qa_k_top_device
    config['model_config']['encoder_config'][
        'k_top_post_selection'] = qa_k_top_post_selection
    config['model_config']['encoder_config'][
        'separate_memory_values'] = separate_memory_values
    if qa_k_top_post_selection is None:
      config['save_k_retrieval'] = None
    self.run_model(config, 1000)
Test loss function runs and produces expected values.
language/mentionmemory/tasks/mention_based_entity_qa_task_test.py