perf improve in python code

get buffer_names_dict only when it's used
microsoft · Nov 6, 2023 · de3c24b · de3c24b
1 parent dfafcb5
commit de3c24b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/orttraining/orttraining/python/training/ortmodule/_io.py b/orttraining/orttraining/python/training/ortmodule/_io.py
@@ -206,7 +206,7 @@ def _expand_inputs(current_input, non_none_inputs, name=""):
     _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 _expand_inputs(current_input, non_none_inputs, name=""):
 
         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: inp for buffer_name, inp in named_buffer}
             try:  # noqa: SIM105
                 inp = buffer_names_dict[name]
             except KeyError: