import torch

from transformers import PreTrainedModel

from configuration_mlp import MLPConfig


class MLP(PreTrainedModel):
    config_class = MLPConfig

    def __init__(self, config):
        super().__init__(config)
        self.input_layer = torch.nn.Linear(config.input_size, config.output_size)

    def forward(self, inputs):
        x = self.input_layer(inputs)
        return x

MLP.register_for_auto_class("AutoModel")