yo

config.json

@@ -5,5 +5,8 @@
   "auto_map": {
     "AutoModel": "custom_model.CustomModel",
     "AutoConfig": "custom_model.CustomModelConfig"
-  }
+  },
+  "architectures": [
+    "CustomLM"
+  ]
 }