| from transformers import PreTrainedModel, BertConfig, AutoModel | |
| class Model(PreTrainedModel): | |
| config_class = BertConfig | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.model = AutoModel.from_pretrained("bert-base-uncased") | |
| def forward(self, **inputs): | |
| outs = self.model(**inputs) | |
| return outs | |