fix

microsoft · Oct 22, 2023 · 02e887b · 02e887b
1 parent d7f9a40
commit 02e887b
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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)