BertConfig typing in modeling_mosaicbert.py

mosaicml · Aug 17, 2023 · 1cd7c2e · 1cd7c2e
1 parent 56aa68c
commit 1cd7c2e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llmfoundry/models/mosaicbert/modeling_mosaicbert.py b/llmfoundry/models/mosaicbert/modeling_mosaicbert.py
@@ -92,7 +92,7 @@ class BertModel(BertPreTrainedModel):
     ```
     """
 
-    def __init__(self, config, add_pooling_layer=True):
+    def __init__(self, config: BertConfig, add_pooling_layer=True):
         super(BertModel, self).__init__(config)
         self.embeddings = BertEmbeddings(config)
         self.encoder = BertEncoder(config)