Fix the PT2E UT (#2071)

Signed-off-by: yiliu30 <yi4.liu@intel.com>

Author: yiliu30

test/3x/torch/quantization/test_pt2e_quant.py

@@ -134,8 +134,6 @@ def calib_fn(model):
         logger.warning("out shape is %s", out.shape)
         assert out is not None
 
-    # TODO: AssertionError: Node 'call_function' <OpOverload(op='quantized_decomposed.quantize_per_tensor', overload='default')> should occur 2 times, but 1
-    @pytest.mark.skipif(True, reason="TODO: fix AssertionError")
     @pytest.mark.skipif(not GT_OR_EQUAL_TORCH_VERSION_2_5, reason="Requires torch>=2.5")
     def test_quantize_simple_model_with_set_local(self, force_not_import_ipex):
         model, example_inputs = self.build_simple_torch_model_and_example_inputs()
@@ -149,12 +147,14 @@ def calib_fn(model):
         quant_config = get_default_static_config()
         quant_config.set_local("fc1", StaticQuantConfig(w_dtype="fp32", act_dtype="fp32"))
         q_model = quantize(model=model, quant_config=quant_config, run_fn=calib_fn)
-
-        # check the half node
         expected_node_occurrence = {
             # Only quantize the `fc2`
-            torch.ops.quantized_decomposed.quantize_per_tensor.default: 2,
-            torch.ops.quantized_decomposed.quantize_per_tensor.default: 2,
+            # Quantize/Dequantize input
+            torch.ops.quantized_decomposed.quantize_per_tensor.default: 1,
+            torch.ops.quantized_decomposed.dequantize_per_tensor.default: 1,
+            # Quantize/Dequantize weight
+            torch.ops.quantized_decomposed.quantize_per_channel.default: 1,
+            torch.ops.quantized_decomposed.quantize_per_channel.default: 1,
         }
         expected_node_occurrence = {
             torch_test_quant_common.NodeSpec.call_function(k): v for k, v in expected_node_occurrence.items()