Update optimum/habana/transformers/models/mt5/modeling_mt5.py

Co-authored-by: Yaser Afshar <yaser.afshar@intel.com>
Gaurav7888 · Aug 5, 2024 · 486eb63 · 486eb63
1 parent 40438c8
commit 486eb63
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/optimum/habana/transformers/models/mt5/modeling_mt5.py b/optimum/habana/transformers/models/mt5/modeling_mt5.py
@@ -80,6 +80,8 @@ def gaudi_MT5Attention_forward(
         """
         Copied from https://github.com/huggingface/transformers/blob/v4.40.2/src/transformers/models/mt5/modeling_mt5.py
         Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
+        The only differences are:
+        - add new args token_idx
         """
         # Input is (batch_size, seq_length, dim)
         # Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length)