Support QDQ GatherElements in quantization tool

microsoft · Aug 12, 2024 · 20be096 · 20be096
1 parent d5264c6
commit 20be096
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/onnxruntime/python/tools/quantization/operators/gather.py b/onnxruntime/python/tools/quantization/operators/gather.py
@@ -55,7 +55,7 @@ def __init__(self, onnx_quantizer, onnx_node):
 
     def quantize(self):
         node = self.node
-        assert node.op_type == "Gather"
+        assert node.op_type == "Gather" or node.op_type == "GatherElements"
 
         if self.quantizer.is_valid_quantize_weight(node.input[0]) or self.quantizer.force_quantize_no_input_check:
             self.quantizer.quantize_activation_tensor(node.input[0])

diff --git a/onnxruntime/python/tools/quantization/registry.py b/onnxruntime/python/tools/quantization/registry.py
@@ -79,6 +79,7 @@
     "MatMul": QDQMatMul,
     "Split": QDQSplit,
     "Gather": QDQGather,
+    "GatherElements": QDQGather,
     "Where": QDQWhere,
     "InstanceNormalization": QDQNormalization,
     "LayerNormalization": QDQNormalization,