[fix] tengine output getitem bug

ModelTC · Sep 21, 2022 · 835b48e · 835b48e
1 parent 8cc2477
commit 835b48e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/mqbench/custom_quantizer/tengine_u8_quantizer.py b/mqbench/custom_quantizer/tengine_u8_quantizer.py
@@ -1,5 +1,6 @@
 import torch
 from torch.fx import GraphModule
+import operator
 
 from mqbench.utils.registry import register_model_quantizer
 from mqbench.utils import getitem2node
@@ -72,6 +73,8 @@ def _find_act_quants(self, model: GraphModule) -> list:
             elif node.op == "output":
                 for _arg in node.args:
                     if isinstance(_arg, torch.fx.node.Node):
+                        if _arg.target == operator.getitem:
+                            continue
                         if _arg.op == 'placeholder':
                             continue
                         node_need_to_quantize_output.append(_arg)