update

config.json/config.json

@@ -1,7 +1,4 @@
 {
-  "auto_map": {
-    "AutoConfig": "model.TestConfig"
-  },
   "input_dim": 4,
   "model_type": "test-model",
   "output_dim": 16,

config.json/model.py

@@ -0,0 +1,21 @@
+from transformers import PreTrainedModel, PretrainedConfig
+from torch import nn
+
+class TestConfig(PretrainedConfig):
+    model_type = "test-model" 
+    
+    def __init__(self, input_dim=4, output_dim=16, **kwargs):
+        super().__init__(**kwargs)
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+    
+    
+class TestModel(PreTrainedModel):
+    config_class = TestConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.layer = nn.Linear(config.input_dim, config.output_dim)
+        
+    def forward(self, input):
+        return self.layer(input)

model.py

@@ -5,9 +5,9 @@ class TestConfig(PretrainedConfig):
     model_type = "test-model" 
     
     def __init__(self, input_dim=4, output_dim=16, **kwargs):
-        super().__init__(**kwargs)
         self.input_dim = input_dim
         self.output_dim = output_dim
+        super().__init__(**kwargs)
     
     
 class TestModel(PreTrainedModel):

pytorch_model.bin/model.py

@@ -0,0 +1,21 @@
+from transformers import PreTrainedModel, PretrainedConfig
+from torch import nn
+
+class TestConfig(PretrainedConfig):
+    model_type = "test-model" 
+    
+    def __init__(self, input_dim=4, output_dim=16, **kwargs):
+        super().__init__(**kwargs)
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+    
+    
+class TestModel(PreTrainedModel):
+    config_class = TestConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.layer = nn.Linear(config.input_dim, config.output_dim)
+        
+    def forward(self, input):
+        return self.layer(input)

pytorch_model.bin/model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:456fdf360a6885691c2e9a7da8262857e3d415c36f57abf4224941b954ec7580
+oid sha256:d70ce77882860863fe8be43f83938f725ebc05732db7d9e06363d74f3dcd4caa
 size 496