def load_weights(cls, config: ml_collections.ConfigDict) -> Dict[str, Any]:
    """Load model weights from file.

    We assume that `encoder_name` is specified in the config.
    We use corresponding class to load encoder weights.

    Args:
      config: experiment config.

    Returns:
      Dictionary of model weights.
    """
    encoder_name = config.model_config.encoder_name
    encoder_class = encoder_registry.get_registered_encoder(encoder_name)
    encoder_variables = encoder_class.load_weights(config)
    model_variables = {}
    for group_key in encoder_variables:
      model_variables[group_key] = {'encoder': encoder_variables[group_key]}

    return model_variables
Load model weights from file.

    We assume that `encoder_name` is specified in the config.
    We use corresponding class to load encoder weights.

    Args:
      config: experiment config.

    Returns:
      Dictionary of model weights.
    
language/mentionmemory/tasks/downstream_encoder_task.py