diff --git a/torch/ao/quantization/pt2e/graph_utils.py b/torch/ao/quantization/pt2e/graph_utils.py index 2d4e73dee3577..2390bafed2868 100644 --- a/torch/ao/quantization/pt2e/graph_utils.py +++ b/torch/ao/quantization/pt2e/graph_utils.py @@ -25,7 +25,7 @@ {torch.nn.BatchNorm2d, torch.nn.functional.batch_norm}, {torch.nn.Hardtanh, torch.nn.functional.hardtanh, torch.nn.functional.hardtanh_}, {torch.add, operator.add, operator.iadd, "add", "add_"}, - {torch.mul, operator.mul, operator.imul}, + {torch.mul, operator.mul, operator.imul, "mul", "mul_"}, ] diff --git a/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py b/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py index 085c65b768bce..e89c06715b2e1 100644 --- a/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py +++ b/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py @@ -739,7 +739,7 @@ def _annotate_mul( filter_fn: Optional[Callable[[Node], bool]] = None, ) -> Optional[List[List[Node]]]: mul_partitions = get_source_partitions( - gm.graph, [operator.mul, torch.mul, operator.imul], filter_fn + gm.graph, ["mul", "mul_", operator.mul, torch.mul, operator.imul], filter_fn ) mul_partitions = list(itertools.chain(*mul_partitions.values())) annotated_partitions = []