From 1f4a3529ddde7796cd018b6c2bd97a6e099c4b88 Mon Sep 17 00:00:00 2001 From: Jambay Kinley Date: Wed, 27 Sep 2023 14:53:37 -0700 Subject: [PATCH] Bugfix: Add initializer to model in AttentionMask directly (#17719) --- .../python/tools/transformers/fusion_attention.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/onnxruntime/python/tools/transformers/fusion_attention.py b/onnxruntime/python/tools/transformers/fusion_attention.py index 40f2aee875..a7460157ba 100644 --- a/onnxruntime/python/tools/transformers/fusion_attention.py +++ b/onnxruntime/python/tools/transformers/fusion_attention.py @@ -78,7 +78,15 @@ class AttentionMask: # ReduceSum-13: axes is moved from attribute to input axes_name = "ort_const_1_reduce_sum_axes" if self.model.get_initializer(axes_name) is None: - self.add_initializer(name=axes_name, data_type=TensorProto.INT64, dims=[1], vals=[1], raw=False) + self.model.add_initializer( + helper.make_tensor( + name=axes_name, + data_type=TensorProto.INT64, + dims=[1], + vals=[1], + raw=False, + ) + ) mask_index_node = helper.make_node( "ReduceSum", inputs=[input_name, axes_name],