Fix

Xilinx · Jul 2, 2023 · 86e7750 · 86e7750
1 parent 2a29be1
commit 86e7750
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/brevitas_examples/llm/llm_quant/mha_layers.py b/src/brevitas_examples/llm/llm_quant/mha_layers.py
@@ -13,6 +13,8 @@ def attention_mask_handler(
             attention_mask = attention_mask.repeat(batch_size, 1, 1, 1)
         if attention_mask.shape[1] == 1:
             attention_mask = attention_mask.repeat(1, num_heads, 1, 1)
+        if attention_mask.shape[2] == 1:
+            attention_mask = attention_mask.repeat(1, 1, query_seq_length, 1)
         attention_mask = attention_mask.view(
             batch_size * num_heads, query_seq_length, key_value_seq_length)
     elif len(attention_mask.shape) == 2 and attention_mask.shape[0] == 1: