Update remaining callsites to not use pt2e quant API from pytorch

Differential Revision: D89744472

Pull Request resolved: #16380

Author: jerryzh168

backends/nxp/quantizer/neutron_quantizer.py

@@ -53,7 +53,6 @@
     no_outside_users,
 )
 from torch import fx
-from torch.ao.quantization.quantizer.utils import _annotate_output_qspec
 from torchao.quantization.pt2e import (
     FakeQuantize,
     FusedMovingAvgObsFakeQuantize,
@@ -62,6 +61,7 @@
     MovingAverageMinMaxObserver,
 )
 from torchao.quantization.pt2e.quantizer import (
+    annotate_output_qspec,
     ComposableQuantizer,
     DerivedQuantizationSpec,
     OperatorConfig,
@@ -338,7 +338,7 @@ def _annotate_inputs(self, model: fx.GraphModule):
                 continue
 
             if node.op == "placeholder" and len(node.users) > 0:
-                _annotate_output_qspec(node, act_qspec(self.is_qat))
+                annotate_output_qspec(node, act_qspec(self.is_qat))
                 self._mark_input_node_as_annotated(node)
 
     def validate(self, model: torch.fx.GraphModule) -> None:

backends/nxp/quantizer/utils.py

@@ -15,13 +15,15 @@
 import torch
 from torch import fx
 from torch._ops import OpOverload
-from torch.ao.quantization import move_exported_model_to_eval
 from torch.export import ExportedProgram
 from torch.fx.passes.utils.source_matcher_utils import (
     check_subgraphs_connected,
     SourcePartition,
 )
-from torchao.quantization.pt2e import ObserverOrFakeQuantize
+from torchao.quantization.pt2e import (
+    move_exported_model_to_eval,
+    ObserverOrFakeQuantize,
+)
 from torchao.quantization.pt2e.quantize_pt2e import (
     convert_pt2e,
     prepare_pt2e,