update

config.json

@@ -1,3 +1,5 @@
 {
-    "hidden_size": 128
-  }
+  "hidden_size": 128,
+  "model_type": "custom-model",
+  "transformers_version": "4.44.2"
+}

custom_model.py

@@ -1,6 +1,5 @@
 # custom_model.py
 from transformers import PreTrainedModel, PretrainedConfig, AutoConfig, AutoModel
-import torch
 import torch.nn as nn
 
 class CustomModelConfig(PretrainedConfig):
@@ -19,6 +18,3 @@ class CustomModel(PreTrainedModel):
     def forward(self, input_ids):
         output = self.linear(input_ids)
         return output
-
-AutoConfig.register("custom-model", CustomModelConfig)
-AutoModel.register(CustomModelConfig, CustomModel)