update

auto.ipynb

@@ -0,0 +1,19 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "558ced93",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

config.json/config.json

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

pytorch_model.bin/model.safetensors

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

train.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import (WEIGHTS_NAME, CONFIG_NAME)
+from transformers import (WEIGHTS_NAME, CONFIG_NAME, AutoConfig)
 
 from model import (TestConfig, TestModel)
 
@@ -14,5 +14,7 @@ if __name__ == "__main__":
     pred = model(x)
     print(pred)
     
+    AutoConfig.register("test-model", TestConfig)
+    
     config.save_pretrained(CONFIG_NAME)
     model.save_pretrained(WEIGHTS_NAME)