Adding model type to the config

configuration_bert.py

@@ -97,6 +97,7 @@ class FlexBertConfig(TransformersBertConfig):
         pad_logits: bool = False,
         compile_model: bool = False,
         masked_prediction: bool = False,
+        model_type: str = "flex_bert"
         **kwargs,
     ):
         """
@@ -213,6 +214,7 @@ class FlexBertConfig(TransformersBertConfig):
         self.pad_logits = pad_logits
         self.compile_model = compile_model
         self.masked_prediction = masked_prediction
+        self.model_type = model_type
 
         if loss_kwargs.get("return_z_loss", False):
             if loss_function != "fa_cross_entropy":