Create __main__.py

linear/__main__.py

@@ -0,0 +1,30 @@
+from transformers import PretrainedConfig, PreTrainedModel
+import torch
+
+
+class LinearConfig(PretrainedConfig):
+    model_type = "linear"
+
+    def __init__(
+        self,
+        in_dims=768,
+        out_dims=1024,
+        **kwargs,
+    ):
+      self.in_dims = in_dims
+      self.out_dims = out_dims
+      super().__init__(**kwargs)
+
+
+class LinearModel(PreTrainedModel):
+    config_class = LinearConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = torch.nn.Linear(
+            config.in_dims,
+            config.out_dims,
+        )
+
+    def forward(self, x):
+        return self.model.forward(x)