update func

huggingface · Jan 22, 2024 · 9aac9c8 · 9aac9c8
1 parent bb7b71a
commit 9aac9c8
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/optimum/onnx/utils.py b/optimum/onnx/utils.py
@@ -16,7 +16,7 @@
 from typing import List, Tuple, Union
 
 import onnx
-from onnx.external_data_helper import ExternalDataInfo, _get_initializer_tensors
+from onnx.external_data_helper import ExternalDataInfo, _get_attribute_tensors, _get_initializer_tensors
 
 
 def _get_onnx_external_data_tensors(model: onnx.ModelProto) -> List[str]:
@@ -30,6 +30,13 @@ def _get_onnx_external_data_tensors(model: onnx.ModelProto) -> List[str]:
         for tensor in model_tensors
         if tensor.HasField("data_location") and tensor.data_location == onnx.TensorProto.EXTERNAL
     ]
+
+    model_tensors = _get_attribute_tensors(model)
+    model_tensors_ext += [
+        ExternalDataInfo(tensor).location
+        for tensor in model_tensors
+        if tensor.HasField("data_location") and tensor.data_location == onnx.TensorProto.EXTERNAL
+    ]
     return model_tensors_ext