Add backward compatible types to pt2e prepare (#2244)

Summary: Pull Request resolved: #2244

Differential Revision: D75248288

Author: metascroy

torchao/quantization/pt2e/quantize_pt2e.py

@@ -25,6 +25,7 @@
     _fuse_conv_bn_,
     _get_node_name_to_scope,
 )
+from typing import Union
 
 from .convert import _convert_to_reference_decomposed_fx
 from .prepare import prepare
@@ -39,7 +40,7 @@
 
 def prepare_pt2e(
     model: GraphModule,
-    quantizer: Quantizer,
+    quantizer: Union[Quantizer, torch.ao.quantization.quantizer.quantizer.Quantizer],
 ) -> GraphModule:
     """Prepare a model for post training quantization
 
@@ -127,7 +128,7 @@ def calibrate(model, data_loader):
 
 def prepare_qat_pt2e(
     model: GraphModule,
-    quantizer: Quantizer,
+    quantizer: Quantizer | torch.ao.quantization.quantizer.quantizer.Quantizer,
 ) -> GraphModule:
     """Prepare a model for quantization aware training