def build_model(cls,
                  model_config: ml_collections.FrozenConfigDict) -> nn.Module:
    """Builds model by instantiating flax module associated with task."""
    return cls.model_class(**model_config)
Builds model by instantiating flax module associated with task.
language/mentionmemory/tasks/base_task.py