From 8d641229e6dbd6364a610923c31fc51448e2601a Mon Sep 17 00:00:00 2001
From: Patrice Vignola <vignola.patrice@gmail.com>
Date: Sun, 10 Dec 2023 21:36:19 -0800
Subject: [PATCH] Fix GQA shape inference (#18723)

The shape inference is always returning before getting the chance to
infer the key/value outputs.
---
 onnxruntime/core/graph/contrib_ops/bert_defs.cc | 1 -
 1 file changed, 1 deletion(-)

diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
index b97fb0d289..ea67218b5c 100644
--- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
@@ -259,7 +259,6 @@ void GroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext&
       *output_shape.add_dim() = query_dims[1];
       *output_shape.add_dim() = query_dims[2];
       updateOutputShape(ctx, 0, output_shape);
-      return;
     } else {
       fail_shape_inference("Missing input 2 (value)");
     }