fix gradient checkpointing in donut_swin

hackyon · Jan 21, 2024 · 1cf345e · 1cf345e
1 parent e834a3e
commit 1cf345e
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/transformers/models/donut/modeling_donut_swin.py b/src/transformers/models/donut/modeling_donut_swin.py
@@ -750,7 +750,12 @@ def forward(
 
             if self.gradient_checkpointing and self.training:
                 layer_outputs = self._gradient_checkpointing_func(
-                    layer_module.__call__, hidden_states, input_dimensions, layer_head_mask, output_attentions
+                    layer_module.__call__,
+                    hidden_states,
+                    input_dimensions,
+                    layer_head_mask,
+                    output_attentions,
+                    always_partition,
                 )
             else:
                 layer_outputs = layer_module(