| import torch | |
| from transformers import BertModel | |
| from .configuration import NewModelConfig | |
| class NewModel(BertModel): | |
| config_class = NewModelConfig | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.last_layer = torch.nn.Linear(config.hidden_size, config.new_hidden_size) |