From 02e887b0a93d5743ded438b8610ef4a89e0a306a Mon Sep 17 00:00:00 2001 From: JiCheng Date: Thu, 19 Oct 2023 06:03:21 +0000 Subject: [PATCH] fix --- onnxruntime/python/tools/transformers/large_model_exporter.py | 1 + 1 file changed, 1 insertion(+) diff --git a/onnxruntime/python/tools/transformers/large_model_exporter.py b/onnxruntime/python/tools/transformers/large_model_exporter.py index ceeb5d218e334..3b344d6dc9342 100644 --- a/onnxruntime/python/tools/transformers/large_model_exporter.py +++ b/onnxruntime/python/tools/transformers/large_model_exporter.py @@ -213,6 +213,7 @@ def fetch_onnx_inputs_outputs_name( ): """fetch onnx inputs and outputs name""" num_of_past_key = 0 + kv_cache_axis = {0: "batch_size"} # try get num_of_past_key and shape of past_key_value if past_key_values is not None: num_of_past_key = len(past_key_values)