code: def prepare_qat_fx(
      model: Union[PreTrainedModel, GraphModule],
      qconfig_dict: Any,
      prepare_custom_config_dict: Optional[Dict[str, Any]] = None,
      backend_config_dict: Optional[Dict[str, Any]] = None,
      input_names: Optional[List[str]] = None,
      check: bool = True,
  ) -> ObservedGraphModule:
      """
      Transformers models compatible version of torch.quantization.quantize_fx.prepare_qat_fx, refer to the [PyTorch
      documentation](https://pytorch.org/docs/stable/generated/torch.quantization.quantize_fx.prepare_qat_fx.html#torch.quantization.quantize_fx.prepare_qat_fx)
      for more details.
  
      Args:
          model (`PreTrainedModel` or `torch.fx.GraphModule`):
              The model to prepare, must be in train mode.
          qconfig_dict (`Any`):
               Dictionary specifying how and which modules and operations should be quantized, refer to the PyTorch
               documentation for more details.
          prepare_custom_config_dict (`Dict[str, Any]`, *optional*):
              Customization configuration dictionary for quantization tool, refer to the PyTorch documentation for more
              details.
          backend_config_dict (`Dict[str, Any]`, *optional*):
              A dictionary that specifies how operators are quantized in a backend, this includes how the operators are
              observed, supported fusion patterns, how quantize/dequantize ops are inserted, supported dtypes etc.
              The structure of the dictionary is still WIP and will change in the future, please don't use right now.
          input_names (`List[str]`, *optional*):
              The input names of the model, only used to trace if model is a PreTrainedModel. This is not needed if model
              is already a GraphModule.
          check (`bool`, *optional*, defaults to `True`):
              If True, a check is done to verify that the model can be traced.
  
      Returns:
          `torch.quantization.fx.graph_module.ObservedGraphModule`: An ObservedGraphModule ready for QAT.
  
      Example:
  
          ```python
          >>> import torch
          >>> from torch.ao.quantization import get_default_qat_qconfig
          >>> from transformers import BertForSequenceClassification
          >>> from optimum.fx.quantization import prepare_qat_fx
  
          >>> model = BertForSequenceClassification.from_pretrained("textattack/bert-base-uncased-SST-2")
          >>> model.train()  # doctest: +IGNORE_RESULT
  
          >>> # Prepare the model
          >>> qconfig = get_default_qat_qconfig('fbgemm')
          >>> qconfig_dict = {"": qconfig}
          >>> prepared_model = prepare_qat_fx(model, qconfig_dict)
          ```
      """
      if check:
          check_if_model_is_supported(model)
      tracer_cls = QuantizationTracer
      if not isinstance(model, GraphModule):
          if input_names is None:
              input_names = model.dummy_inputs.keys()
          input_names = list(input_names)
          tracer_cls = specialized_quantization_tracer_creator(
              get_concrete_args(model, input_names)
          )
      orig_quantization_tracer = (
          torch.ao.quantization.quantize_fx.QuantizationTracer
      )
      torch.ao.quantization.quantize_fx.QuantizationTracer = tracer_cls
      graph_module = orig_prepare_qat_fx(
          model,
          qconfig_dict,
          prepare_custom_config_dict=prepare_custom_config_dict,
          backend_config_dict=backend_config_dict,
      )
      torch.ao.quantization.quantize_fx.QuantizationTracer = (
          orig_quantization_tracer
      )
      return graph_module
  
imports: from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type, Union
  import torch
  from torch.fx.node import Argument, Node, Target
  from torch.nn.intrinsic import _FusedModule
  from torch.quantization.fx.graph_module import GraphModule, ObservedGraphModule
  from torch.quantization.quantize_fx import Scope, ScopeContextManager
  from torch.quantization.quantize_fx import fuse_fx
  from torch.quantization.quantize_fx import prepare_fx
  from torch.quantization.quantize_fx import prepare_qat_fx
  from transformers import PreTrainedModel
  from transformers.utils.fx import HFTracer, check_if_model_is_supported, get_concrete_args, symbolic_trace
  from ..utils import check_if_available
label_desc: Use image node_id 4 for calling a global function with example usage: prepare_qat_fx(model, qconfig_dict, prepare_custom_config_dict, backend_config_dict, input_names, check) and returns: graph_module