| 90 | self.load_state_dict(checkpoint) |
| 91 | |
| 92 | def init_weights(self): |
| 93 | for module in self.decoder.modules(): |
| 94 | if isinstance(module, (nn.Linear, nn.Embedding)): |
| 95 | module.weight.data.normal_(mean=0.0, std=0.02) |
| 96 | elif isinstance(module, nn.LayerNorm): |
| 97 | module.bias.data.zero_() |
| 98 | module.weight.data.fill_(1.0) |
| 99 | if isinstance(module, nn.Linear) and module.bias is not None: |
| 100 | module.bias.data.zero_() |
| 101 | for p in self.generator.parameters(): |
| 102 | if p.dim() > 1: |
| 103 | xavier_uniform_(p) |
| 104 | else: |
| 105 | p.data.zero_() |
| 106 | |
| 107 | def forward( |
| 108 | self, encoder_input_ids, decoder_input_ids, token_type_ids, encoder_attention_mask, decoder_attention_mask, |