awesome-model_v3 / model.py
valhalla's picture
add model
fbf3c3e
raw
history blame contribute delete
393 Bytes
import torch
from transformers import PreTrainedModel
from .config import AwesomeConfig
class AwesomeModel(PreTrainedModel):
config_class = AwesomeConfig
base_model_prefix = "base"
def __init__(self, config):
super().__init__(config)
self.linear = torch.nn.Linear(config.hidden_size, config.hidden_size)
def forward(self, x):
return self.linear(x)