add additional args to custom_prepare_pt2e

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/654

example_inputs and is_qat args are needed for some models during prepare_pt2e step.

Reviewed By: chakriu, tarun292

Differential Revision: D54873270

fbshipit-source-id: 67df457aca82fd9da77969133ecf390cdc80fb85

d2go/quantization/modeling.py

@@ -362,7 +362,9 @@ def prepare_fake_quant_model(cfg, model, is_qat, example_input=None):
     """
     if cfg.QUANTIZATION.PT2E:  # pt2e quantization
         if hasattr(model, "custom_prepare_pt2e"):
-            model, convert_pt2e_callback = model.custom_prepare_pt2e(cfg)
+            model, convert_pt2e_callback = model.custom_prepare_pt2e(
+                cfg, is_qat, example_input
+            )
         else:
             logger.info("Using default pt2e quantization APIs with XNNPACKQuantizer")
             captured_model = capture_pre_autograd_graph(model, example_input)