FLAN-T5 has bad performance when using regional compilation (#1744)

huggingface · Feb 5, 2025 · 9b59b75 · 9b59b75
1 parent d15132c
commit 9b59b75
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/habana/transformers/models/t5/modeling_t5.py b/optimum/habana/transformers/models/t5/modeling_t5.py
@@ -406,7 +406,7 @@ def gaudi_T5Stack_forward(
 
         if self.gradient_checkpointing and self.training:
             layer_outputs = self._gradient_checkpointing_func(
-                layer_module.forward,
+                layer_module.__call__,
                 hidden_states,
                 extended_attention_mask,
                 position_bias,