| 392 | |
| 393 | |
| 394 | class BertEncoder(nn.Module): |
| 395 | def __init__(self, config): |
| 396 | super().__init__() |
| 397 | self.config = config |
| 398 | self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)]) |
| 399 | |
| 400 | def forward( |
| 401 | self, |
| 402 | hidden_states, |
| 403 | attention_mask=None, |
| 404 | head_mask=None, |
| 405 | encoder_hidden_states=None, |
| 406 | encoder_attention_mask=None, |
| 407 | output_attentions=False, |
| 408 | output_hidden_states=False, |
| 409 | ): |
| 410 | all_hidden_states = () |
| 411 | all_attentions = () |
| 412 | for i, layer_module in enumerate(self.layer): |
| 413 | if output_hidden_states: |
| 414 | all_hidden_states = all_hidden_states + (hidden_states,) |
| 415 | |
| 416 | if getattr(self.config, "gradient_checkpointing", False): |
| 417 | |
| 418 | def create_custom_forward(module): |
| 419 | def custom_forward(*inputs): |
| 420 | return module(*inputs, output_attentions) |
| 421 | |
| 422 | return custom_forward |
| 423 | |
| 424 | layer_outputs = torch.utils.checkpoint.checkpoint( |
| 425 | create_custom_forward(layer_module), |
| 426 | hidden_states, |
| 427 | attention_mask, |
| 428 | head_mask[i], |
| 429 | encoder_hidden_states, |
| 430 | encoder_attention_mask, |
| 431 | ) |
| 432 | else: |
| 433 | layer_outputs = layer_module( |
| 434 | hidden_states, |
| 435 | attention_mask, |
| 436 | head_mask[i], |
| 437 | encoder_hidden_states, |
| 438 | encoder_attention_mask, |
| 439 | output_attentions, |
| 440 | ) |
| 441 | hidden_states = layer_outputs[0] |
| 442 | |
| 443 | if output_attentions: |
| 444 | all_attentions = all_attentions + (layer_outputs[1],) |
| 445 | |
| 446 | # Add last layer |
| 447 | if output_hidden_states: |
| 448 | all_hidden_states = all_hidden_states + (hidden_states,) |
| 449 | |
| 450 | outputs = (hidden_states,) |
| 451 | if output_hidden_states: |