From 630c877b43bb6f13271afb52e8fd5e008314028d Mon Sep 17 00:00:00 2001
From: zhijiang <43435212+zhijxu-MS@users.noreply.github.com>
Date: Tue, 7 Nov 2023 15:24:57 +0800
Subject: [PATCH] Zhijxu/improve ortmodule python perf a little bit (#13716)

improve 2 python functions a little bit.

according to a profiling result from a real user case, we find that 2
python function can be improved. the first is the result before
improvement, the second is after improvement, we can see 8ms saved from
the improvement.

![image](https://user-images.githubusercontent.com/43435212/202961725-b88d679e-993b-4910-a339-253f3ed5dcde.png)

![image](https://user-images.githubusercontent.com/43435212/202961732-6c6deebf-962f-4392-90d7-03705433e3ee.png)
---
 orttraining/orttraining/python/training/ortmodule/_io.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/orttraining/orttraining/python/training/ortmodule/_io.py b/orttraining/orttraining/python/training/ortmodule/_io.py
index e7c1b30daa..1b6e2df9d2 100644
--- a/orttraining/orttraining/python/training/ortmodule/_io.py
+++ b/orttraining/orttraining/python/training/ortmodule/_io.py
@@ -206,7 +206,7 @@ def _combine_input_buffers_initializers(
     _expand_inputs(inputs, non_none_inputs)
     flattened_kwargs_inputs = {}
     _expand_inputs(kwargs, flattened_kwargs_inputs)
-    buffer_names_dict = {buffer_name: inp for buffer_name, inp in named_buffer}
+    buffer_names_dict = None
     result = []
     embed_sparsity_results = OrderedDict()
     label_sparsity_results = OrderedDict()
@@ -232,6 +232,8 @@ def _combine_input_buffers_initializers(
 
         if inp is None:
             # Registered buffers are translated to user_input+initializer in ONNX
+            if buffer_names_dict is None:
+                buffer_names_dict = {buffer_name: i for buffer_name, i in named_buffer}
             try:  # noqa: SIM105
                 inp = buffer_names_dict[name]
             except KeyError: