Initialize the weights.
(self, module)
| 305 | super().__init__(*inputs, **kwargs) |
| 306 | |
| 307 | def _init_weights(self, module): |
| 308 | """Initialize the weights.""" |
| 309 | if isinstance(module, (nn.Linear,)): |
| 310 | # Slightly different from Mesh Transformer JAX which uses truncated_normal for initialization |
| 311 | # cf https://github.com/pytorch/pytorch/pull/5617 |
| 312 | module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) |
| 313 | if module.bias is not None: |
| 314 | module.bias.data.zero_() |
| 315 | elif isinstance(module, nn.Embedding): |
| 316 | module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) |
| 317 | if module.padding_idx is not None: |
| 318 | module.weight.data[module.padding_idx].zero_() |
| 319 | elif isinstance(module, nn.LayerNorm): |
| 320 | module.bias.data.zero_() |
| 321 | module.weight.data.fill_(1.0) |
| 322 | |
| 323 | def _set_gradient_checkpointing(self, module, value=False): |
| 324 | if isinstance(module, MossModel): |
nothing calls this directly
no outgoing calls
no test coverage detected