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