From 9ef9c39728854f06d66db781dccc583f70c90fcd Mon Sep 17 00:00:00 2001 From: Lysandre Debut Date: Mon, 7 Sep 2020 10:56:08 +0200 Subject: [PATCH] Cannot index `None` (#6984) --- src/transformers/modeling_bert.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/transformers/modeling_bert.py b/src/transformers/modeling_bert.py index 1cebc53240..e3f9cd72de 100755 --- a/src/transformers/modeling_bert.py +++ b/src/transformers/modeling_bert.py @@ -464,6 +464,8 @@ class BertEncoder(nn.Module): if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) + layer_head_mask = head_mask[i] if head_mask is not None else None + if getattr(self.config, "gradient_checkpointing", False): def create_custom_forward(module): @@ -476,7 +478,7 @@ class BertEncoder(nn.Module): create_custom_forward(layer_module), hidden_states, attention_mask, - head_mask[i], + layer_head_mask, encoder_hidden_states, encoder_attention_mask, ) @@ -484,7 +486,7 @@ class BertEncoder(nn.Module): layer_outputs = layer_module( hidden_states, attention_mask, - head_mask[i], + layer_head_mask, encoder_hidden_states, encoder_attention_mask, output_attentions,