Update bert_layers.py

bert_layers.py

@@ -51,6 +51,7 @@ from transformers.models.bert.modeling_bert import BertPreTrainedModel
 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
@@ -625,6 +626,8 @@ class BertModel(BertPreTrainedModel):
     ```
     """
 
+    config_class = BertConfig
+
     def __init__(self, config, add_pooling_layer=True):
         super(BertModel, self).__init__(config)
         self.embeddings = BertEmbeddings(config)
@@ -758,6 +761,8 @@ class BertLMHeadModel(BertPreTrainedModel):
 
 class BertForMaskedLM(BertPreTrainedModel):
 
+    config_class = BertConfig
+
     def __init__(self, config):
         super().__init__(config)
 
@@ -928,6 +933,8 @@ class BertForSequenceClassification(BertPreTrainedModel):
     e.g., GLUE tasks.
     """
 
+    config_class = BertConfig
+
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels