Error with model_type

configuration_vatrpp.py

@@ -1,5 +1,6 @@
-from transformers import PretrainedConfig
+from transformers import PretrainedConfig, AutoConfig
 
+@AutoConfig.register("vatrpp")
 class VATrPPConfig(PretrainedConfig):
     model_type = "vatrpp"
 

modeling_vatrpp.py

@@ -1,4 +1,4 @@
-from transformers import PreTrainedModel
+from transformers import PreTrainedModel, AutoModel
 from .configuration_vatrpp import VATrPPConfig
 import json
 import os
@@ -20,7 +20,6 @@ from util.misc import FakeArgs
 from util.text import TextGenerator
 from util.vision import detect_text_bounds
 from torchvision.transforms.functional import to_pil_image
-from transformers import CONFIG_MAPPING, MODEL_MAPPING
 
 
 def get_long_tail_chars():
@@ -31,6 +30,7 @@ def get_long_tail_chars():
 
     return chars
 
+@AutoModel.register(VATrPPConfig)
 class VATrPP(PreTrainedModel):
     config_class = VATrPPConfig
 
@@ -130,7 +130,3 @@ class VATrPP(PreTrainedModel):
             x_pos += word.shape[1] + gap_width
 
         return result
-
-
-CONFIG_MAPPING.register("vatrpp", VATrPPConfig)
-MODEL_MAPPING.register(VATrPPConfig, VATrPP)