workaround for torch.utils.checkpoint.checkpoint not accepting bool

huggingface · Jun 12, 2020 · 5eb68bb · 5eb68bb
1 parent b36648f
commit 5eb68bb
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/src/transformers/modeling_bert.py b/src/transformers/modeling_bert.py
@@ -413,14 +413,20 @@ def forward(
                 all_hidden_states = all_hidden_states + (hidden_states,)
 
             if getattr(self.config, "gradient_checkpointing", False):
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, output_attentions)
+
+                    return custom_forward
+
                 layer_outputs = torch.utils.checkpoint.checkpoint(
-                    layer_module,
+                    create_custom_forward(layer_module),
                     hidden_states,
                     attention_mask,
                     head_mask[i],
                     encoder_hidden_states,
                     encoder_attention_mask,
-                    output_attentions,
                 )
             else:
                 layer_outputs = layer_module(