Update modeling_norbert.py

modeling_norbert.py

@@ -414,10 +414,11 @@ class Classifier(nn.Module):
             nn.Dropout(drop_out),
             nn.Linear(config.hidden_size, num_labels)
         )
-        self.initialize(config.hidden_size)
+        self.hidden_size = config.hidden_size
+        self.initialize()
 
-    def initialize(self, hidden_size):
-        std = math.sqrt(2.0 / (5.0 * hidden_size))
+    def initialize(self):
+        std = math.sqrt(2.0 / (5.0 * self.hidden_size))
         nn.init.trunc_normal_(self.nonlinearity[1].weight, mean=0.0, std=std, a=-2*std, b=2*std)
         nn.init.trunc_normal_(self.nonlinearity[-1].weight, mean=0.0, std=std, a=-2*std, b=2*std)
         self.nonlinearity[1].bias.data.zero_()
@@ -438,6 +439,9 @@ class NorbertForSequenceClassification(NorbertModel):
         self.num_labels = config.num_labels
         self.head = Classifier(config, self.num_labels)
 
+    def post_init(self):
+        self.head.initialize()
+
     def forward(
         self,
         input_ids: Optional[torch.Tensor] = None,
@@ -504,6 +508,9 @@ class NorbertForTokenClassification(NorbertModel):
         self.num_labels = config.num_labels
         self.head = Classifier(config, self.num_labels)
 
+    def post_init(self):
+        self.head.initialize()
+
     def forward(
         self,
         input_ids: Optional[torch.Tensor] = None,