Update modeling_norbert.py

modeling_norbert.py

@@ -277,12 +277,12 @@ class NorbertPreTrainedModel(PreTrainedModel):
 
 
 class NorbertModel(NorbertPreTrainedModel):
-    def __init__(self, config, add_mlm_layer=False):
-        super().__init__(config)
+    def __init__(self, config, add_mlm_layer=False, gradient_checkpointing=False, **kwargs):
+        super().__init__(config, **kwargs)
         self.config = config
 
         self.embedding = Embedding(config)
-        self.transformer = Encoder(config, activation_checkpointing=False)
+        self.transformer = Encoder(config, activation_checkpointing=gradient_checkpointing)
         self.classifier = MaskClassifier(config, self.embedding.word_embedding.weight) if add_mlm_layer else None
 
     def get_input_embeddings(self):
@@ -352,8 +352,8 @@ class NorbertModel(NorbertPreTrainedModel):
 class NorbertForMaskedLM(NorbertModel):
     _keys_to_ignore_on_load_unexpected = ["head"]
 
-    def __init__(self, config):
-        super().__init__(config, add_mlm_layer=True)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, add_mlm_layer=True, **kwargs)
 
     def get_output_embeddings(self):
         return self.classifier.nonlinearity[-1].weight
@@ -432,8 +432,8 @@ class NorbertForSequenceClassification(NorbertModel):
     _keys_to_ignore_on_load_unexpected = ["classifier"]
     _keys_to_ignore_on_load_missing = ["head"]
 
-    def __init__(self, config):
-        super().__init__(config, add_mlm_layer=False)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = config.num_labels
         self.head = Classifier(config, self.num_labels)
@@ -498,8 +498,8 @@ class NorbertForTokenClassification(NorbertModel):
     _keys_to_ignore_on_load_unexpected = ["classifier"]
     _keys_to_ignore_on_load_missing = ["head"]
 
-    def __init__(self, config):
-        super().__init__(config, add_mlm_layer=False)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = config.num_labels
         self.head = Classifier(config, self.num_labels)
@@ -546,8 +546,8 @@ class NorbertForQuestionAnswering(NorbertModel):
     _keys_to_ignore_on_load_unexpected = ["classifier"]
     _keys_to_ignore_on_load_missing = ["head"]
 
-    def __init__(self, config):
-        super().__init__(config, add_mlm_layer=False)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = config.num_labels
         self.head = Classifier(config, self.num_labels)
@@ -614,8 +614,8 @@ class NorbertForMultipleChoice(NorbertModel):
     _keys_to_ignore_on_load_unexpected = ["classifier"]
     _keys_to_ignore_on_load_missing = ["head"]
 
-    def __init__(self, config):
-        super().__init__(config, add_mlm_layer=False)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = getattr(config, "num_labels", 2)
         self.head = Classifier(config, self.num_labels)