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