Update modeling_norbert.py

modeling_norbert.py

@@ -231,15 +231,12 @@ class NorbertPreTrainedModel(PreTrainedModel):
     def _init_weights(self, module):
         std = math.sqrt(2.0 / (5.0 * self.hidden_size))
 
-        if isinstance(module, nn.Linear):
+        if isinstance(module, nn.Linear) or isinstance(module, nn.Embedding):
             nn.init.trunc_normal_(module.weight.data, mean=0.0, std=std, a=-2*std, b=2*std)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            nn.init.trunc_normal_(module.weight.data, mean=0.0, std=std, a=-2*std, b=2*std)
-        elif isinstance(module, nn.LayerNorm):
-            module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm) and module.weight is not None:
             module.weight.data.fill_(1.0)
+        if module.bias is not None:
+            module.bias.data.zero_()
 
 
 class NorbertModel(NorbertPreTrainedModel):