| import torch | |
| from transformers import PreTrainedModel | |
| from hg.config import MConfig | |
| class MModel(PreTrainedModel): | |
| config_class = MConfig | |
| def __init__(self, config: MConfig): | |
| super().__init__(config) | |
| self.m = torch.nn.Sequential( | |
| torch.nn.Linear(1, config.hidden), | |
| torch.nn.ReLU(), | |
| torch.nn.Linear(config.hidden, 1), | |
| ) | |
| self.post_init() | |
| def forward(self, input: torch.Tensor, label: torch.Tensor): | |
| return torch.nn.functional.mse_loss(self.m(input), label) | |