diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py
index b1686f7418..a37c48f259 100644
--- a/onnxscript/function_libs/torch_lib/ops/nn.py
+++ b/onnxscript/function_libs/torch_lib/ops/nn.py
@@ -1835,7 +1835,7 @@ def aten_scaled_dot_product_attention(
     if is_causal:
         attn_mask = _causal_attention_mask(query, key)
 
-    if enable_gqa:
+    if enable_gqa and query.shape[1] != key.shape[1]:
         key, value = _attention_repeat_kv_for_group_query(query, key, value)
 
     if attn_mask is None: