fix GQA error message

Signed-off-by: Charlene Yang <[email protected]>
NVIDIA · Nov 12, 2024 · 062a7d0 · 062a7d0
1 parent 237b493
commit 062a7d0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
@@ -7951,7 +7951,7 @@ def forward(
             assert (
                 key_layer.shape[-2] == self.num_gqa_groups_per_partition
                 and value_layer.shape[-2] == self.num_gqa_groups_per_partition
-            ), f"Keys and values must have num_gqa_group = {self.num_gqa_groups} heads!"
+            ), f"Keys and values must have num_gqa_group = {self.num_gqa_groups_per_partition} heads!"
             assert qkv_format in [
                 "sbhd",
                 "bshd",