From 062a7d07421be42ceed86b9e6be7451f65b96b66 Mon Sep 17 00:00:00 2001
From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Date: Tue, 12 Nov 2024 10:19:34 -0800
Subject: [PATCH] fix GQA error message

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---
 transformer_engine/pytorch/attention.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
index 6b153fd3c1..7a401629f3 100644
--- 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",