def load_weights(config: ml_collections.ConfigDict) -> Dict[str, Any]:
    """Load model weights and mention memory."""

    if config.load_weights == 'memory_only':
      model_variables = {}
    else:
      model_variables = base_encoder.BaseEncoder.load_weights(config)
    memory_variables = MentionMemoryEncoder.load_memory(config)
    model_variables['constants'] = memory_variables

    return model_variables
Load model weights and mention memory.
language/mentionmemory/encoders/mention_memory_encoder.py