update

config.json/config.json

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

config.json/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/config.json

@@ -2,6 +2,10 @@
   "architectures": [
     "TestModel"
   ],
+  "auto_map": {
+    "AutoConfig": "model.TestConfig",
+    "AutoModel": "model.TestModel"
+  },
   "input_dim": 4,
   "model_type": "test-model",
   "output_dim": 16,

pytorch_model.bin/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.safetensors

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