def METHOD_NAME(self) -> Union[nn.Module, TorchModel]:
    """ Instantiate a pytorch model and return.
    By default, we will create a model using config from configuration file. You can
    override this method in a subclass.
    """
    model_args = {} if self.num_labels is None else {
        'num_labels': self.num_labels
    }
    model = Model.from_pretrained(
        self.model_dir, cfg_dict=self.cfg, **model_args)
    if not isinstance(model, nn.Module) and hasattr(model, 'model'):
        return model.model
    elif isinstance(model, nn.Module):
        return model