Upload model

config.json

@@ -1,5 +1,12 @@
 {
+  "architectures": [
+    "FCModel"
+  ],
+  "auto_map": {
+    "AutoModel": "model_fc.FCModel"
+  },
   "model_type": "fc_layer",
   "num_nodes": 10,
+  "torch_dtype": "float32",
   "transformers_version": "4.28.1"
 }

configuration_fc.py

@@ -0,0 +1,12 @@
+from transformers import PretrainedConfig
+from typing import List
+
+class FCConfig(PretrainedConfig):
+    model_type = "fc_layer" # For registration auto class
+
+    def __init__(self, num_nodes=10, **kwargs):
+        ## Must include keyword arguments
+        ## Note: Must need a default value for hyperparameters
+        self.num_nodes = num_nodes
+        super().__init__(**kwargs)
+

model_fc.py

@@ -0,0 +1,15 @@
+from transformers import PreTrainedModel
+from .configuration_fc import FCConfig
+from torch.nn import Linear
+
+class FCModel(PreTrainedModel):
+    config_class = FCConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Linear(in_features=10, out_features=config.num_nodes)
+    
+    def forward(self, tensor):
+        # Use as forward similar to forward in torch
+        return self.model.forward(tensor)
+