add autoconfig

bert_layers.py

@@ -25,6 +25,9 @@ from .bert_padding import (index_first_axis,
                                             index_put_first_axis, pad_input,
                                             unpad_input, unpad_input_only)
 
+from .configuration_bert import BertConfig
+
+
 try:
     from .flash_attn_triton import flash_attn_qkvpacked_func
 except ImportError as e:
@@ -683,6 +686,8 @@ class BertOnlyNSPHead(nn.Module):
 
 
 class BertForMaskedLM(BertPreTrainedModel):
+    
+    config_class = BertConfig
 
     def __init__(self, config):
         super().__init__(config)
@@ -812,6 +817,9 @@ class BertForMaskedLM(BertPreTrainedModel):
 
 
 class BertForSequenceClassification(BertPreTrainedModel):
+    
+    config_class = BertConfig
+    
     """Bert Model transformer with a sequence classification/regression head.
 
     This head is just a linear layer on top of the pooled output. Used for,