Address comments

microsoft · Mar 21, 2024 · 174f3c0 · 174f3c0
1 parent 1f8b9a8
commit 174f3c0
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 5 deletions.
diff --git a/onnxruntime/python/tools/transformers/optimizer.py b/onnxruntime/python/tools/transformers/optimizer.py
@@ -42,7 +42,7 @@
 from onnx_model_vae import VaeOnnxModel
 
 import onnxruntime
-from onnxruntime.python.tools.transformers.optimizer_utils import extract_external_data_from_model
+from onnxruntime.transformers.optimizer_utils import extract_external_data_from_model
 
 logger = logging.getLogger(__name__)
 

diff --git a/onnxruntime/python/tools/transformers/optimizer_utils.py b/onnxruntime/python/tools/transformers/optimizer_utils.py
@@ -1,7 +1,13 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
 from fusion_utils import NumpyHelper
-from onnxruntime import OrtValue
-from onnx.external_data_helper import set_external_data
 from onnx import ModelProto
+from onnx.external_data_helper import set_external_data
+
+from onnxruntime import OrtValue
+
 
 def extract_external_data_from_model(model: ModelProto):
     """

diff --git a/onnxruntime/test/python/transformers/test_optimizer_utils.py b/onnxruntime/test/python/transformers/test_optimizer_utils.py
@@ -1,7 +1,13 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
 import unittest
+
 import numpy
 from onnx import ModelProto, TensorProto, helper
-from onnxruntime.python.tools.transformers.optimizer_utils import extract_external_data_from_model
+
+from onnxruntime.transformers.optimizer_utils import extract_external_data_from_model
 
 
 class TestOptimizerUtils(unittest.TestCase):
@@ -12,7 +18,6 @@ def test_extract_external_data_from_model(self):
         self.assertEqual(len(external_values), 1)
         self.assertEqual(external_values[0].numpy(), [0.0])
 
-
     def _get_model_proto_with_raw_data(self) -> ModelProto:
         input = helper.make_tensor_value_info("inputs", TensorProto.FLOAT, [None])
         output = helper.make_tensor_value_info("outputs", TensorProto.FLOAT, [None])