updates

config.json

@@ -4,10 +4,10 @@
     "FlexBertModel",
     "FlexBertForCasualLM"
   ],
-  "auto_map": {                                                                                                                                                                                                               
-    "AutoConfig": "orionweller/test-flex-gpt--configuration_bert.FlexBertConfig",
-    "AutoModel": "orionweller//test-flex-gpt--modeling_flexbert.FlexBertModel",
-    "AutoModelForCasualLM": "orionweller//test-flex-gpt--modeling_flexbert.FlexBertForCasualLM"                                                                                                                                                                                                                
+  "auto_map": {
+      "AutoConfig": "orionweller/test-flex-gpt--configuration_bert.FlexBertConfig",
+      "AutoModel": "orionweller/test-flex-gpt--modeling_flexbert.FlexBertModel",
+      "AutoModelForCausalLM": "orionweller/test-flex-gpt--modeling_flexbert.FlexBertForCasualLM"
   },
   "attention_layer": "rope",
   "attention_probs_dropout_prob": 0.0,

modeling_flexbert.py

@@ -1734,6 +1734,7 @@ class FlexBertForCasualLM(FlexBertPreTrainedModel):
         params += _count_parameters(self.lm_head, trainable)
         return params
 
+FlexBertForCasualLM.register_for_auto_class("AutoModelForCausalLM")
 
 def init_model_from_pretrained(
     pretrained_model: FlexBertModel,