Repository: quic/aimet (language: python)

Issue:
[PART 4] Deprecate pytorch autoquant v1



Discussion / hints:

Fix patch:
diff --git a/Examples/torch/quantization/autoquant_v2.ipynb b/Examples/torch/quantization/autoquant.ipynb
similarity index 97%
rename from Examples/torch/quantization/autoquant_v2.ipynb
rename to Examples/torch/quantization/autoquant.ipynb
index 3135679943..e936943280 100644
--- a/Examples/torch/quantization/autoquant_v2.ipynb
+++ b/Examples/torch/quantization/autoquant.ipynb
@@ -24,12 +24,7 @@
     "3. Run AutoQuant\n",
     "\n",
     "#### What this notebook is not\n",
-    "This notebook is not designed to show state-of-the-art AutoQuant results. For example, it uses a relatively quantization-friendly model like Resnet18. Also, some optimization parameters are deliberately chosen to have the notebook execute more quickly.\n",
-    "\n",
-    "<div class=\"alert alert-info\">\n",
-    "    NOTE: This notebook is for auto_quant_v2.AutoQuant. For examples of auto_quant.AutoQuant (will be deprecated), see \n",
-    "    <a href=\"autoquant_v1.ipynb\">autoquant_v1.ipynb</a>.\n",
-    "</div>"
+    "This notebook is not designed to show state-of-the-art AutoQuant results. For example, it uses a relatively quantization-friendly model like Resnet18. Also, some optimization parameters are deliberately chosen to have the notebook execute more quickly.\n"
    ]
   },
   {
@@ -186,7 +181,7 @@
    },
    "outputs": [],
    "source": [
-    "from aimet_torch.auto_quant_v2 import AutoQuant\n",
+    "from aimet_torch.auto_quant import AutoQuant\n",
     "\n",
     "class UnlabeledDatasetWrapper(Dataset):\n",
     "    def __init__(self, dataset):\n",
diff --git a/Examples/torch/quantization/autoquant_v1.ipynb b/Examples/torch/quantization/autoquant_v1.ipynb
deleted file mode 100644
index a9ee37d711..0000000000
--- a/Examples/torch/quantization/autoquant_v1.ipynb
+++ /dev/null
@@ -1,332 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
-   "source": [
-    "# AutoQuant\n",
-    "\n",
-    "This notebook shows a working code example of how to use AIMET AutoQuant feature.\n",
-    "\n",
-    "AIMET offers a suite of neural network post-training quantization (PTQ) techniques that can be applied in succession. However, the process of finding the right combination and sequence of techniques to apply is time-consuming and requires careful analysis, which can be challenging especially for non-expert users. We instead recommend AutoQuant to save time and effort.\n",
-    "\n",
-    "AutoQuant is an API that applies various PTQ techniques in AIMET automatically based on analyzing the model and best-known heuristics. In AutoQuant, users specify the amount of tolerable accuracy drop, and AutoQuant will apply PTQ techniques cumulatively until the target accuracy is satisfied.\n",
-    "\n",
-    "\n",
-    "#### Overall flow\n",
-    "This notebook covers the following\n",
-    "1. Define constants and helper functions\n",
-    "2. Load a pretrained FP32 model\n",
-    "3. Run AutoQuant\n",
-    "\n",
-    "#### What this notebook is not\n",
-    "This notebook is not designed to show state-of-the-art AutoQuant results. For example, it uses a relatively quantization-friendly model like Resnet18. Also, some optimization parameters are deliberately chosen to have the notebook execute more quickly.\n",
-    "\n",
-    "<div class=\"alert alert-info\">\n",
-    "    NOTE: This notebook is for auto_quant.AutoQuant (will be deprecated). For examples of auto_quant_v2.AutoQuant, see \n",
-    "    <a href=\"autoquant_v2.ipynb\">autoquant_v2.ipynb</a>.\n",
-    "</div>"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "---\n",
-    "## Dataset\n",
-    "\n",
-    "This notebook relies on the ImageNet dataset for the task of image classification. If you already have a version of the dataset readily available, please use that. Else, please download the dataset from appropriate location (e.g. https://image-net.org/challenges/LSVRC/2012/index.php#).\n",
-    "\n",
-    "**Note1**: The ImageNet dataset typically has the following characteristics and the dataloader provided in this example notebook rely on these\n",
-    "- Subfolders 'train' for the training samples and 'val' for the validation samples. Please see the [pytorch dataset description](https://pytorch.org/vision/0.8/_modules/torchvision/datasets/imagenet.html) for more details.\n",
-    "- A subdirectory per class, and a file per each image sample\n",
-    "\n",
-    "**Note2**: To speed up the execution of this notebook, you may use a reduced subset of the ImageNet dataset. E.g. the entire ILSVRC2012 dataset has 1000 classes, 1000 training samples per class and 50 validation samples per class. But for the purpose of running this notebook, you could perhaps reduce the dataset to say 2 samples per class. This exercise is left upto the reader and is not necessary.\n",
-    "\n",
-    "Edit the cell below and specify the directory where the downloaded ImageNet dataset is saved."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "is_executing": true
-    }
-   },
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "from torchvision import transforms, datasets\n",
-    "\n",
-    "DATASET_DIR = '/path/to/dataset'   # Please replace this with a real directory\n",
-    "\n",
-    "val_transforms = transforms.Compose([\n",
-    "    transforms.CenterCrop(224),\n",
-    "    transforms.ToTensor(),\n",
-    "    transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n",
-    "])\n",
-    "\n",
-    "imagenet_dataset = datasets.ImageFolder(root=os.path.join(DATASET_DIR, 'val'), transform=val_transforms)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 1. Define Constants and Helper functions\n",
-    "\n",
-    "In this section the constants and helper functions needed to run this eaxmple are defined.\n",
-    "\n",
-    "- **EVAL_DATASET_SIZE** A typical value is 5000. In this notebook, this value has been set to 500 for faster execution.\n",
-    "- **CALIBRATION_DATASET_SIZE** A typical value is 2000. In this notebook, this value has been set to 200 for faster execution.\n",
-    "\n",
-    "\n",
-    "The helper function **_create_sampled_data_loader()** returns a DataLoader based on the dataset and the number of samples provided."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "import random\n",
-    "from typing import Optional\n",
-    "from tqdm import tqdm\n",
-    "import torch\n",
-    "from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler, Subset\n",
-    "from aimet_torch.utils import in_eval_mode, get_device\n",
-    "\n",
-    "EVAL_DATASET_SIZE = 500\n",
-    "CALIBRATION_DATASET_SIZE = 200\n",
-    "\n",
-    "_datasets = {}\n",
-    "\n",
-    "def _create_sampled_data_loader(dataset, num_samples):\n",
-    "    if num_samples not in _datasets:\n",
-    "        indices = random.sample(range(len(dataset)), num_samples)\n",
-    "        _datasets[num_samples] = Subset(dataset, indices)\n",
-    "    return DataLoader(_datasets[num_samples], batch_size=32)\n",
-    "\n",
-    "\n",
-    "def eval_callback(model: torch.nn.Module, num_samples: Optional[int] = None) -> float:\n",
-    "    if num_samples is None:\n",
-    "        num_samples = EVAL_DATASET_SIZE\n",
-    "\n",
-    "    data_loader = _create_sampled_data_loader(imagenet_dataset, num_samples)\n",
-    "    device = get_device(model)\n",
-    "    \n",
-    "    correct = 0\n",
-    "    with in_eval_mode(model), torch.no_grad():\n",
-    "        for image, label in tqdm(data_loader):\n",
-    "            image = image.to(device)\n",
-    "            label = label.to(device)\n",
-    "            logits = model(image)\n",
-    "            top1 = logits.topk(k=1).indices\n",
-    "            correct += (top1 == label.view_as(top1)).sum()\n",
-    "\n",
-    "    return int(correct) / num_samples"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 2. Load a pretrained FP32 model\n",
-    "For this example, we are going to load a pretrained resnet18 model from torchvision. Similarly, you can load any pretrained PyTorch model instead."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "from torchvision.models import resnet18\n",
-    "\n",
-    "model = resnet18(pretrained=True).eval()\n",
-    "\n",
-    "if torch.cuda.is_available():\n",
-    "    model.to(torch.device('cuda'))\n",
-    "\n",
-    "accuracy = eval_callback(model)\n",
-    "print(f'- FP32 accuracy: {accuracy}')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 3. Run AutoQuant\n",
-    "\n",
-    "### Prepare and Validate Model\n",
-    "AIMET quantization features require the user's model definition to follow certain guidelines. For example, functionals defined in forward pass should be changed to equivalent torch.nn.Module. AIMET user guide lists all these guidelines. The following ModelPreparer API uses new graph transformation feature available in PyTorch 1.9+ version and automates model definition changes required to comply with the above guidelines."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from aimet_torch.model_preparer import prepare_model\n",
-    "from aimet_torch.model_validator.model_validator import ModelValidator\n",
-    "\n",
-    "model = prepare_model(model)\n",
-    "dummy_input = torch.randn((1, 3, 224, 224))\n",
-    "result = ModelValidator.validate_model(model, model_input=dummy_input.to(get_device(model)))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Create AutoQuant Object\n",
-    "\n",
-    "The AutoQuant feature utilizes an unlabeled dataset to achieve quantization. The class **UnlabeledDatasetWrapper** creates an unlabeled Dataset object from a labeled Dataset.\n",
-    "\n",
-    "The **allowed_accuracy_drop** parameter is set by the user to convey to the AutoQuant feature, how much accuracy drop is tolerated by the user. AutoQuant applies a series of quantization features. When the allowed accuracy is reached, AutoQuant stops applying any subsequent quantization feature. Please refer AutoQuant User Guide and API documentation for complete details."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "from aimet_torch.auto_quant import AutoQuant\n",
-    "\n",
-    "class UnlabeledDatasetWrapper(Dataset):\n",
-    "    def __init__(self, dataset):\n",
-    "        self._dataset = dataset\n",
-    "\n",
-    "    def __len__(self):\n",
-    "        return len(self._dataset)\n",
-    "\n",
-    "    def __getitem__(self, index):\n",
-    "        images, _ = self._dataset[index]\n",
-    "        return images\n",
-    "\n",
-    "\n",
-    "unlabeled_imagenet_dataset = UnlabeledDatasetWrapper(imagenet_dataset)\n",
-    "unlabeled_imagenet_data_loader = _create_sampled_data_loader(unlabeled_imagenet_dataset,\n",
-    "                                                             CALIBRATION_DATASET_SIZE)\n",
-    "\n",
-    "auto_quant = AutoQuant(allowed_accuracy_drop=0.01,\n",
-    "                       unlabeled_dataset_iterable=unlabeled_imagenet_data_loader,\n",
-    "                       eval_callback=eval_callback)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Set AdaRound Parameters (optional)\n",
-    "The AutoQuant feature internally uses default parameters to execute the AdaRound step.\n",
-    "If and only if necessary, the default AdaRound Parameters should be modified using the API shown below.\n",
-    "\n",
-    "**Note:**\n",
-    "To execute this example faster, the default value of the **num_iterations** parameter has been reduced from 10000 to 2000"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "from aimet_torch.adaround.adaround_weight import AdaroundParameters\n",
-    "\n",
-    "ADAROUND_DATASET_SIZE = 200\n",
-    "adaround_data_loader = _create_sampled_data_loader(unlabeled_imagenet_dataset, ADAROUND_DATASET_SIZE)\n",
-    "adaround_params = AdaroundParameters(adaround_data_loader, num_batches=len(adaround_data_loader), default_num_iterations=2000)\n",
-    "auto_quant.set_adaround_params(adaround_params)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Run AutoQuant\n",
-    "\n",
-    "This step applies the AutoQuant feature. The best possible quantized model, the associated eval_score and the path to the best encoding files are returned."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "scrolled": false
-   },
-   "outputs": [],
-   "source": [
-    "model.eval()\n",
-    "model, accuracy, encoding_path =\\\n",
-    "    auto_quant.apply(model.cuda(),\n",
-    "                     dummy_input_on_cpu=dummy_input.cpu(),\n",
-    "                     dummy_input_on_gpu=dummy_input.cuda())\n",
-    "print(f\"- Quantized Accuracy (after optimization):  {accuracy}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
-   "source": [
-    "---\n",
-    "## Summary\n",
-    "\n",
-    "Hope this notebook was useful for you to understand how to use AIMET AutoQuant feature.\n",
-    "\n",
-    "Few additional resources\n",
-    "- Refer to the AIMET API docs to know more details of the APIs and parameters\n",
-    "- Refer to the other example notebooks to understand how to use AIMET CLE and AdaRound features in a standalone fashion."
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.8.0"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/auto_quant_v2.py b/TrainingExtensions/torch/src/python/aimet_torch/auto_quant_v2.py
deleted file mode 100644
index 65c6cf69c7..0000000000
--- a/TrainingExtensions/torch/src/python/aimet_torch/auto_quant_v2.py
+++ /dev/null
@@ -1,1341 +0,0 @@
-# -*- mode: python -*-
-# =============================================================================
-#  @@-COPYRIGHT-START-@@
-#
-#  Copyright (c) 2022, Qualcomm Innovation Center, Inc. All rights reserved.
-#
-#  Redistribution and use in source and binary forms, with or without
-#  modification, are permitted provided that the following conditions are met:
-#
-#  1. Redistributions of source code must retain the above copyright notice,
-#     this list of conditions and the following disclaimer.
-#
-#  2. Redistributions in binary form must reproduce the above copyright notice,
-#     this list of conditions and the following disclaimer in the documentation
-#     and/or other materials provided with the distribution.
-#
-#  3. Neither the name of the copyright holder nor the names of its contributors
-#     may be used to endorse or promote products derived from this software
-#     without specific prior written permission.
-#
-#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
-#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-#  POSSIBILITY OF SUCH DAMAGE.
-#
-#  SPDX-License-Identifier: BSD-3-Clause
-#
-#  @@-COPYRIGHT-END-@@
-# =============================================================================
-# pylint: disable=too-many-lines
-
-"""Temporary buffer file for adding new features to AutoQuant"""
-import copy
-import contextlib
-from collections import OrderedDict, defaultdict
-from dataclasses import dataclass
-import functools
-import itertools
-import math
-import traceback
-import os
-import sys
-import io
-from unittest.mock import patch
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union, Mapping
-import pickle
-from uuid import uuid4
-import torch
-from torch.utils.data import DataLoader
-import jinja2
-from bokeh.resources import CDN
-from tqdm import tqdm
-
-from aimet_torch import utils
-from aimet_torch.adaround.adaround_weight import Adaround, AdaroundParameters
-from aimet_torch.cross_layer_equalization import equalize_model
-from aimet_torch.batch_norm_fold import fold_all_batch_norms
-from aimet_torch.quantsim import QuantizationSimModel
-from aimet_torch.utils import get_all_quantizers, in_eval_mode
-from aimet_torch.onnx_utils import OnnxExportApiArgs
-from aimet_torch.model_preparer import prepare_model
-from aimet_torch.model_validator.model_validator import ModelValidator
-
-from aimet_common.auto_quant import Diagnostics
-from aimet_common.cache import Cache
-from aimet_common.defs import QuantScheme
-from aimet_common.utils import AimetLogger, Spinner
-from aimet_common.quantsim import validate_quantsim_inputs
-
-
-_logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.AutoQuant)
-
-cache = Cache()
-
-
-# The number of samples to be used for performance evaluation.
-# NOTE: None means "all".
-NUM_SAMPLES_FOR_PERFORMANCE_EVALUATION = None
-
-
-@dataclass(frozen=True)
-class _QuantSchemePair:
-    param_quant_scheme: QuantScheme
-    output_quant_scheme: QuantScheme
-    param_percentile: Optional[float] = None
-    output_percentile: Optional[float] = None
-
-    def __str__(self):
-        def scheme_to_str(quant_scheme, percentile):
-            if quant_scheme == QuantScheme.post_training_percentile:
-                return f"{percentile}%ile"
-            if quant_scheme in (QuantScheme.post_training_tf,
-                                QuantScheme.training_range_learning_with_tf_init):
-                return "tf"
-            if quant_scheme in (QuantScheme.post_training_tf_enhanced,
-                                QuantScheme.training_range_learning_with_tf_enhanced_init):
-                return "tf-enhanced"
-            raise ValueError
-
-        param_str = scheme_to_str(self.param_quant_scheme, self.param_percentile)
-        output_str = scheme_to_str(self.output_quant_scheme, self.output_percentile)
-        return f"W@{param_str} / A@{output_str}"
-
-
-_QUANT_SCHEME_CANDIDATES = (
-    # Weight:     tf
-    # Activation: tf
-    _QuantSchemePair(QuantScheme.post_training_tf,
-                     QuantScheme.post_training_tf),
-
-    # Weight:     tf_enhanced
-    # Activation: tf
-    _QuantSchemePair(QuantScheme.post_training_tf_enhanced,
-                     QuantScheme.post_training_tf),
-
-    # Weight:     tf_enhanced
-    # Activation: tf_enhanced
-    _QuantSchemePair(QuantScheme.post_training_tf_enhanced,
-                     QuantScheme.post_training_tf_enhanced),
-
-    # Weight:     tf_enhanced
-    # Activation: percentile(99.9)
-    _QuantSchemePair(QuantScheme.post_training_tf_enhanced,
-                     QuantScheme.post_training_percentile,
-                     output_percentile=99.9),
-
-    # Weight:     tf_enhanced
-    # Activation: percentile(99.99)
-    _QuantSchemePair(QuantScheme.post_training_tf_enhanced,
-                     QuantScheme.post_training_percentile,
-                     output_percentile=99.99),
-)
-
-
-def _validate_inputs(model: torch.nn.Module, # pylint: disable=too-many-arguments
-                     data_loader: DataLoader,
-                     eval_callback: Callable[[torch.nn.Module], float],
-                     dummy_input: torch.Tensor,
-                     results_dir: str,
-                     strict_validation: bool,
-                     quant_scheme: QuantScheme,
-                     param_bw: int,
-                     output_bw: int,
-                     rounding_mode: str):
-    """
-    Confirms inputs are of the correct type
-    :param model: Model to be quantized
-    :param data_loader: A collection that iterates over an unlabeled dataset, used for computing encodings
-    :param eval_callback: Function that calculates the evaluation score
-    :param dummy_input: Dummy input for the model
-    :param results_dir: Directory to save the results of PTQ techniques
-    :param strict_validation: Flag set to True by default. When False, AutoQuant will proceed with execution and try to handle errors internally if possible. This may produce unideal or unintuitive results.
-    :param quant_scheme: Quantization scheme
-    :param param_bw: Parameter bitwidth
-    :param output_bw: Output bitwidth
-    :param rounding_mode: Rounding mode
-    """
-    if not isinstance(model, torch.nn.Module):
-        raise ValueError('Model must be of type torch.nn.Module, not ' + str(type(model).__name__))
-
-    if not isinstance(data_loader, DataLoader):
-        raise ValueError('data_loader must be of type DataLoader, not ' + str(
-            type(data_loader).__name__))
-
-    if not isinstance(eval_callback, Callable):  # pylint: disable=isinstance-second-argument-not-valid-type
-        raise ValueError('eval_callback must be of type Callable, not ' + str(type(eval_callback).__name__))
-
-    if not isinstance(dummy_input, (torch.Tensor, Tuple)):
-        raise ValueError(
-            'dummy_input must be of type torch.Tensor or Tuple, not ' + str(type(dummy_input).__name__))
-
-    if not isinstance(results_dir, str):
-        raise ValueError('results_dir must be of type str, not ' + str(type(results_dir).__name__))
-
-    results_dir = os.path.abspath(results_dir)
-    os.makedirs(results_dir, exist_ok=True)
-
-    if not isinstance(strict_validation, bool):
-        raise ValueError('strict_validation must be of type bool, not ' + str(type(strict_validation).__name__))
-
-    validate_quantsim_inputs(quant_scheme, rounding_mode, output_bw, param_bw)
-
-
-class AutoQuant: # pylint: disable=too-many-instance-attributes
-    """
-    Integrate and apply post-training quantization techniques.
-
-    AutoQuant includes 1) batchnorm folding, 2) cross-layer equalization,
-    and 3) Adaround.
-    These techniques will be applied in a best-effort manner until the model
-    meets the evaluation goal given as allowed_accuracy_drop.
-    """
-
-    def __init__( # pylint: disable=too-many-arguments, too-many-locals
-            self,
-            model: torch.nn.Module,
-            dummy_input: Union[torch.Tensor, Tuple],
-            data_loader: DataLoader,
-            eval_callback: Callable[[torch.nn.Module], float],
-            param_bw: int = 8,
-            output_bw: int = 8,
-            quant_scheme: QuantScheme = QuantScheme.post_training_tf_enhanced,
-            rounding_mode: str = 'nearest',
-            config_file: str = None,
-            results_dir: str = "/tmp",
-            cache_id: str = None,
-            strict_validation: bool = True) -> None:
-        '''
-        :param model: Model to be quantized. Assumes model is on the correct device
-        :param dummy_input: Dummy input for the model. Assumes that dummy_input is on the correct device
-        :param data_loader: A collection that iterates over an unlabeled dataset, used for computing encodings
-        :param eval_callback: Function that calculates the evaluation score
-        :param param_bw: Parameter bitwidth
-        :param output_bw: Output bitwidth
-        :param quant_scheme: Quantization scheme
-        :param rounding_mode: Rounding mode
-        :param config_file: Path to configuration file for model quantizers
-        :param results_dir: Directory to save the results of PTQ techniques
-        :param cache_id: ID associated with cache results
-        :param strict_validation: Flag set to True by default.hen False, AutoQuant will proceed with execution and handle errors internally if possible. This may produce unideal or unintuitive results.
-        '''
-        _validate_inputs(model, data_loader, eval_callback, dummy_input, results_dir,
-                         strict_validation, quant_scheme, param_bw, output_bw, rounding_mode)
-
-        self.fp32_model = model
-        self.dummy_input = dummy_input
-        self.data_loader = data_loader
-        self.eval_callback = eval_callback
-
-        self._quantsim_params = dict(
-            param_bw=param_bw,
-            output_bw=output_bw,
-            quant_scheme=_QuantSchemePair(quant_scheme, quant_scheme),
-            rounding_mode=rounding_mode,
-            config_file=config_file,
-        )
-
-        self.results_dir = results_dir
-        if cache_id:
-            self.cache_dir = os.path.join(results_dir, ".auto_quant_cache", cache_id)
-        else:
-            self.cache_dir = None
-
-        def forward_pass_callback(model, _: Any = None):
-            device = utils.get_device(model)
-            with in_eval_mode(model), torch.no_grad():
-                for input_data in tqdm(data_loader):
-                    input_data = utils.change_tensor_device_placement(input_data, device)
-                    if isinstance(input_data, torch.Tensor):
-                        model(input_data)
-                    else:
-                        assert isinstance(input_data, (tuple, list))
-                        model(*input_data)
-
-        self.forward_pass_callback = forward_pass_callback
-
-        @functools.wraps(eval_callback)
-        def eval_callback_wrapper(model: torch.nn.Module, *args, **kwargs) -> float:
-            """
-            Wrapper to ensure that model is in eval mode before entering eval_callback.
-            """
-            with in_eval_mode(model), torch.no_grad():
-                return eval_callback(model, *args, **kwargs)
-
-        self.eval_callback = eval_callback_wrapper
-
-        # Use at most 2000 samples for AdaRound.
-        num_samples = min(len(self.data_loader.dataset), 2000)
-        batch_size = self.data_loader.batch_size or 1
-        num_batches = math.ceil(num_samples / batch_size)
-        num_batches = min(num_batches, len(self.data_loader))
-        self.adaround_params = AdaroundParameters(self.data_loader, num_batches)
-
-        self._export_kwargs = dict(
-            onnx_export_args=OnnxExportApiArgs(),
-            propagate_encodings=False,
-        )
-        self._model_preparer_kwargs = dict(
-            modules_to_exclude=None,
-            module_classes_to_exclude=None,
-            concrete_args=None,
-        )
-
-        self.eval_manager = _EvalManager(
-            quantsim_factory=self._create_quantsim_and_encodings,
-            eval_func=self._evaluate_model_performance,
-            dummy_input_on_cpu=utils.change_tensor_device_placement(dummy_input, torch.device("cpu")),
-            results_dir=self.results_dir,
-            strict_validation=strict_validation)
-
-        self._quant_scheme_candidates = _QUANT_SCHEME_CANDIDATES
-        self._fp32_acc = None
-
-    def _evaluate_model_performance(self, model) -> float:
-        """
-        Evaluate the model performance.
-        """
-        return self.eval_callback(model, NUM_SAMPLES_FOR_PERFORMANCE_EVALUATION)
-
-    def run_inference(self) -> Tuple[QuantizationSimModel, float]:
-        '''
-        Creates a quantization model and performs inference
-
-        :return: QuantizationSimModel, model accuracy as float
-        '''
-        model = self.fp32_model
-
-        with self.eval_manager.session("Prepare Model") as sess:
-            model = sess.wrap(self._prepare_model)(self.fp32_model)
-
-        # Batchnorm Folding
-        with self.eval_manager.session("Batchnorm Folding", ptq=True) as sess:
-            model, _ = sess.wrap(self._apply_batchnorm_folding)(model)
-            if sess.ptq_result is None:
-                sess.set_ptq_result(model=model,
-                                    applied_techniques=["batchnorm_folding"],
-                                    export_kwargs=self._export_kwargs)
-
-        sim = self._create_quantsim_and_encodings(model)
-
-        if sess.ptq_result is None:
-            # BN folding failed. Need to measure the eval score
-            acc = self._evaluate_model_performance(sim.model)
-        else:
-            # BN folding success. No need to measure the eval score again
-            acc = sess.ptq_result.accuracy
-
-        return sim, acc
-
-    def optimize(self, allowed_accuracy_drop: float = 0.0) -> Tuple[torch.nn.Module, float, str]:
-        """
-        Integrate and apply post-training quantization techniques.
-
-        :param allowed_accuracy_drop: Maximum allowed accuracy drop
-        :return: Tuple of (best model, eval score, encoding path)
-        """
-        result = self._optimize_helper(self._optimize_main, allowed_accuracy_drop)
-        return result["model"],\
-               result["accuracy"],\
-               result["encoding_path"]
-
-    def set_adaround_params(self, adaround_params: AdaroundParameters) -> None:
-        """
-        Set Adaround parameters.
-        If this method is not called explicitly by the user, AutoQuant will use
-        `data_loader` (passed to `__init__`) for Adaround.
-
-        :param adaround_params: Adaround parameters.
-        """
-        self.adaround_params = adaround_params
-
-    def set_export_params(self,
-                          onnx_export_args: OnnxExportApiArgs = -1,
-                          propagate_encodings: bool = None) -> None:
-        """
-        Set parameters for QuantizationSimModel.export.
-
-        :param onnx_export_args: optional export argument with onnx specific overrides
-                if not provide export via torchscript graph
-        :param propagate_encodings: If True, encoding entries for intermediate ops
-                (when one PyTorch ops results in multiple ONNX nodes) are filled with
-                the same BW and data_type as the output tensor for that series of ops.
-        """
-        # Here, we use -1 to indicate `onnx_export_args` wasn't specified
-        # since onnx_export_args being None has its own meaning.
-        if onnx_export_args != -1:
-            self._export_kwargs.update(onnx_export_args=onnx_export_args)
-        if propagate_encodings is not None:
-            self._export_kwargs.update(propagate_encodings=propagate_encodings)
-
-    def set_model_preparer_params(
-            self,
-            modules_to_exclude: Optional[List[torch.nn.Module]] = None,
-            module_classes_to_exclude: Optional[List[torch.nn.Module]] = None,
-            concrete_args: Optional[Dict[str, Any]] = None,
-    ):
-        """
-        Set parameters for model preparer.
-
-        :param modules_to_exclude: List of modules to exclude when tracing.
-        :param module_classes_to_exclude: List of module classes to exclude when tracing.
-        :param concrete_args: Parameter for model preparer. Allows you to partially specialize
-            your function, whether it's to remove control flow or data structures. If the
-            model has control flow, torch.fx won't be able to trace the model. Check
-            torch.fx.symbolic_trace API in detail.
-        """
-        self._model_preparer_kwargs["modules_to_exclude"] = copy.copy(modules_to_exclude)
-        self._model_preparer_kwargs["module_classes_to_exclude"] = copy.copy(module_classes_to_exclude)
-        self._model_preparer_kwargs["concrete_args"] = copy.copy(concrete_args)
-
-    def _create_quantsim_and_encodings( # pylint: disable=too-many-arguments, too-many-locals, too-many-branches
-            self,
-            model: torch.nn.Module,
-            rounding_mode: str = None,
-            output_bw: int = None,
-            output_quant_scheme: QuantScheme = None,
-            output_percentile: float = None,
-            param_bw: int = None,
-            param_quant_scheme: QuantScheme = None,
-            param_percentile: float = None,
-            config_file: str = None,
-            encoding_path: str = None,
-    ) -> QuantizationSimModel:
-        """
-        Create a QuantizationSimModel and compute encoding. If `encoding_path` is not None,
-        it is prioritized over other arguments (`output_bw`, `param_bw`, ...).
-
-        :param model: Model to quantize.
-        :param rounding_mode: Rounding mode. Defaults to self._quantsim_params["rounding_mode"].
-        :param output_bw: Default bitwidth (4-31) to use for quantizing layer inputs andoutputs.
-            Defaults to self._quantsim_params["output_bw"].
-        :param output_quant_scheme: Quantization scheme for output quantizers.
-            Defaults to self._quantsim_params["quant_scheme"].output_quant_scheme.
-        :param output_percentile: Percentile value for outputs.
-            Only valid if output quant scheme is percentile scheme.
-        :param param_bw: Default bitwidth (4-31) to use for quantizing layer parameters.
-            Defaults to self._quantsim_params["param_bw"].
-        :param param_quant_scheme: Quantization scheme for param quantizers.
-            Defaults to self._quantsim_params["quant_scheme"].param_quant_scheme.
-        :param param_percentile: Percentile value for parameters.
-            Only valid if param quant scheme is percentile scheme.
-        :param config_file: Path to configuration file for model quantizers.
-                            Defaults to self._quantsim_params["config_file"].
-        :param encoding_path: Path to parameter encodings file.
-        :return: Quantsim model.
-        """
-        if output_bw is not None:
-            assert output_bw <= 32
-
-        if param_bw is not None:
-            assert param_bw <= 32
-
-        if output_quant_scheme is None or param_quant_scheme is None:
-            assert self._quantsim_params["quant_scheme"] is not None
-
-        kwargs = dict(
-            rounding_mode=(rounding_mode or self._quantsim_params["rounding_mode"]),
-            default_output_bw=(output_bw or self._quantsim_params["output_bw"]),
-            default_param_bw=(param_bw or self._quantsim_params["param_bw"]),
-            config_file=(config_file or self._quantsim_params["config_file"]),
-        )
-        sim = QuantizationSimModel(model, self.dummy_input, **kwargs)
-
-        param_quantizers, input_quantizers, output_quantizers = utils.get_all_quantizers(sim.model)
-
-        default_quant_scheme = self._quantsim_params.get("quant_scheme")
-        if default_quant_scheme is not None:
-            output_quant_scheme = output_quant_scheme or\
-                                  default_quant_scheme.output_quant_scheme
-            output_percentile = output_percentile or default_quant_scheme.output_percentile
-            param_quant_scheme = param_quant_scheme or\
-                                 default_quant_scheme.param_quant_scheme
-            param_percentile = param_percentile or default_quant_scheme.param_percentile
-
-        # Set input/output quantizers' quant schemes
-        for quantizer in itertools.chain(input_quantizers, output_quantizers):
-            quantizer.quant_scheme = output_quant_scheme
-            if quantizer.quant_scheme == QuantScheme.post_training_percentile and\
-                    output_percentile is not None:
-                quantizer.set_percentile_value(output_percentile)
-
-        # Set param quantizers' quant schemes
-        for quantizer in param_quantizers:
-            quantizer.quant_scheme = param_quant_scheme or\
-                                     default_quant_scheme.param_quant_scheme
-            if quantizer.quant_scheme == QuantScheme.post_training_percentile and\
-                    param_percentile is not None:
-                quantizer.set_percentile_value(param_percentile)
-
-        if encoding_path:
-            sim.set_and_freeze_param_encodings(encoding_path)
-
-        param_quantizers, input_quantizers, output_quantizers = utils.get_all_quantizers(sim.model)
-
-        # Disable input/output quantizers, using fp32 to simulate int32.
-        if output_bw == 32:
-            for quantizer in input_quantizers + output_quantizers:
-                quantizer.enabled = False
-
-        # Disable param quantizers, using fp32 to simulate int32.
-        if param_bw == 32:
-            for quantizer in param_quantizers:
-                quantizer.enabled = False
-
-        # Skip encoding computation if none of the quantizers are enabled
-        if any(quantizer.enabled for quantizer in param_quantizers +\
-                                                  input_quantizers +\
-                                                  output_quantizers):
-            sim.compute_encodings(self.forward_pass_callback, None)
-
-        return sim
-
-    def _prepare_model(self, model):
-        prepared_model = prepare_model(model, **self._model_preparer_kwargs)
-
-        if ModelValidator.validate_model(prepared_model, self.dummy_input):
-            _logger.info(
-                "Model validation has succeeded. Proceeding to AutoQuant algorithm."
-            )
-        else:
-            raise ValueError(
-                "Model validation has failed."
-                " Please make the necessary changes to the model and run again."
-            )
-        return prepared_model
-
-    @cache.mark("batchnorm_folding")
-    def _apply_batchnorm_folding(self, model: torch.nn.Module)\
-            -> Tuple[torch.nn.Module, List[Tuple]]:
-        """
-        Apply batchnorm folding.
-
-        NOTE: Input model is not mutated.
-
-        :param model: Model to apply batchnorm folding.
-        :return: Output model and folded pairs.
-        """
-        model = copy.deepcopy(model)
-        folded_pairs = fold_all_batch_norms(model, None, self.dummy_input)
-        return model, folded_pairs
-
-    @cache.mark("cle")
-    def _apply_cross_layer_equalization(self, model: torch.nn.Module) -> torch.nn.Module:
-        """
-        Apply cross-layer equalization.
-
-        NOTE: Input model is not mutated.
-
-        :param model: Model to apply cross-layer-equalization.
-        :return: Output model.
-        """
-        model = copy.deepcopy(model)
-        if isinstance(self.dummy_input, torch.Tensor):
-            input_shape = tuple(self.dummy_input.shape)
-        else:
-            input_shape = [tuple(x.shape) for x in self.dummy_input]
-        equalize_model(model, input_shape)
-        return model
-
-    @cache.mark("adaround")
-    def _apply_adaround(self, model: torch.nn.Module) -> Tuple[torch.nn.Module, str]:
-        """
-        Apply adaround.
-
-        NOTE1: Input model is not mutated.
-        NOTE2: Parameters `param_bw_override_list` and `ignore_quant_ops_list` are always set to None.
-
-        :param model: Model to apply adaround.
-        :return: Output model and the path to the parameter encoding file.
-        """
-        # NOTE: We dont need to make a deepcopy of model here, since Adaround.apply_adaround
-        # internally creates and returns a deepcopy of model.
-        filename_prefix = "adaround"
-        adaround_encoding_path = os.path.join(self.results_dir,
-                                              "{}.encodings".format(filename_prefix))
-
-        sim = self._create_quantsim_and_encodings(model)
-
-        _, input_quantizers, output_quantizers = get_all_quantizers(sim.model)
-        for quantizer in itertools.chain(input_quantizers, output_quantizers):
-            quantizer.enabled = False
-
-        model = Adaround._apply_adaround(sim, model, self.dummy_input, self.adaround_params, # pylint: disable=protected-access
-                                         path=self.results_dir, filename_prefix=filename_prefix)
-
-        return model, adaround_encoding_path
-
-    def _optimize_helper(
-            self,
-            optimize_fn: Callable,
-            allowed_accuracy_drop: float) -> Tuple[torch.nn.Module, float, str]:
-        """
-        Integrate and apply post-training quantization techniques.
-
-        :param allowed_accuracy_drop: Maximum allowed accuracy drop
-        :return: Tuple of (best model, eval score, encoding path)
-        """
-        allowed_accuracy_drop = float(allowed_accuracy_drop)
-        if allowed_accuracy_drop < 0:
-            raise ValueError(
-                "`allowed_accuracy_drop` must be a positive value. Got {:.2f}"
-                .format(allowed_accuracy_drop)
-            )
-
-        self.eval_manager.clear()
-
-        try:
-            with in_eval_mode(self.fp32_model), cache.enable(self.cache_dir):
-                _logger.info("Starting AutoQuant")
-
-                self._fp32_acc = self._evaluate_model_performance(self.fp32_model)
-                target_acc = self._fp32_acc - allowed_accuracy_drop
-                _logger.info("Target eval score: %f", target_acc)
-                _logger.info("FP32 eval score (W32A32): %f", self._fp32_acc)
-
-                ret = optimize_fn(self.fp32_model, target_acc)
-
-                acc = ret["accuracy"]
-                if acc is not None:
-                    _logger.info("Best eval score: %f", acc)
-
-                    if acc < target_acc:
-                        _logger.info(
-                            "AutoQuant is unable to match the target accuracy. "
-                            "Consider Quantization Aware Training."
-                        )
-
-                return ret
-        finally:
-            self.eval_manager.export_diagnostics()
-
-    def get_quant_scheme_candidates(self) -> Tuple[_QuantSchemePair, ...]:
-        """
-        Return the candidates for quant scheme search.
-        During :meth:`~AutoQuant.optimize`, the candidate with the highest accuracy
-        will be selected among them.
-
-        :return: Candidates for quant scheme search
-        """
-        return self._quant_scheme_candidates
-
-    def set_quant_scheme_candidates(self, candidates: Tuple[_QuantSchemePair, ...]):
-        """
-        Set candidates for quant scheme search.
-        During :meth:`~AutoQuant.optimize`, the candidate with the highest accuracy
-        will be selected among them.
-
-        :param candidates: Candidates for quant scheme search
-        """
-        self._quant_scheme_candidates = copy.copy(candidates)
-
-    def _choose_default_quant_scheme(self):
-        def eval_fn(pair: _QuantSchemePair):
-            sim = self._create_quantsim_and_encodings(
-                self.fp32_model,
-                param_quant_scheme=pair.param_quant_scheme,
-                param_percentile=pair.param_percentile,
-                output_quant_scheme=pair.output_quant_scheme,
-                output_percentile=pair.output_percentile,
-            )
-            eval_score = self._evaluate_model_performance(sim.model)
-            _logger.info("Evaluation finished: %s (eval score: %f)", pair, eval_score)
-            return eval_score
-
-        param_bw = self._quantsim_params["param_bw"]
-        output_bw = self._quantsim_params["output_bw"]
-
-        candidates = self.get_quant_scheme_candidates()
-
-        # If the weight representation has sufficient precision (i.e. bitwidth >= 16),
-        # always use tf scheme
-        if param_bw >= 16:
-            candidates = [
-                candidate for candidate in candidates
-                if candidate.param_quant_scheme == QuantScheme.post_training_tf
-            ]
-
-        # If the output representation has sufficient precision (i.e. bitwidth >= 16),
-        # always use tf scheme
-        if output_bw >= 16:
-            candidates = [
-                candidate for candidate in candidates
-                if candidate.output_quant_scheme == QuantScheme.post_training_tf
-            ]
-
-        # If we have only one candidate left, we don't need to evaluated
-        # the quant scheme for comparison
-        if len(candidates) == 1:
-            return candidates[0]
-
-        assert candidates
-
-        # Find the quant scheme that yields the best eval score
-        return max(candidates, key=eval_fn)
-
-    def _optimize_main(self, fp32_model: torch.nn.Module, target_acc: float):
-        """
-        Helper function of apply().
-
-        :param fp32_model: Model to apply PTQ techniques.
-        :param target_acc: Target eval score.
-
-        :raises RuntimeError: If none of the PTQ techniques were finished successfully.
-
-        :return: The best ptq result as a dictionary.
-        """
-        with self.eval_manager.session("Prepare Model") as sess:
-            fp32_model = sess.wrap(self._prepare_model)(self.fp32_model)
-
-        # Choose quant scheme automatically.
-        with self.eval_manager.session("QuantScheme Selection") as sess:
-            self._quantsim_params["quant_scheme"] = sess.wrap(self._choose_default_quant_scheme)()
-
-        with self.eval_manager.session("W32 Evaluation") as sess:
-            w32_eval_score = sess.wrap(sess.eval)(model=fp32_model, param_bw=32)
-            _logger.info("Evaluation finished: W32A%d (eval score: %f)",
-                         self._quantsim_params["output_bw"], w32_eval_score)
-
-            # Early exit
-            if w32_eval_score < target_acc:
-                _logger.info(
-                    "W32A%d eval score (%f) is lower "
-                    "than the target eval score (%f). This means it is unlikely that "
-                    "the target eval score can be met using PTQ techniques. "
-                    "Please consider finetuning the model using range learning.",
-                    self._quantsim_params["output_bw"], w32_eval_score, target_acc
-                )
-
-                # Since AutoQuant pipeline exited early, all the return values are set to None
-                return {
-                    "model": None,
-                    "accuracy": None,
-                    "encoding_path": None,
-                    "applied_techniques": None,
-                }
-
-            sess.result["target_satisfied"] = True
-
-        # Batchnorm Folding
-        with self.eval_manager.session("Batchnorm Folding", ptq=True) as sess:
-            model, _ = sess.wrap(self._apply_batchnorm_folding)(fp32_model)
-            if sess.ptq_result is None:
-                sess.set_ptq_result(model=model,
-                                    applied_techniques=["batchnorm_folding"],
-                                    export_kwargs=self._export_kwargs)
-
-        best_result = self.eval_manager.get_best_ptq_result()
-        if best_result and best_result.accuracy >= target_acc:
-            sess.result["target_satisfied"] = True
-            return best_result.as_dict()
-
-        # Cross-Layer Equalization
-        with self.eval_manager.session("Cross-Layer Equalization", ptq=True) as sess:
-            model = sess.wrap(self._apply_cross_layer_equalization)(fp32_model)
-            if sess.ptq_result is None:
-                sess.set_ptq_result(model=model,
-                                    applied_techniques=["cross_layer_equalization"],
-                                    export_kwargs=self._export_kwargs)
-
-        best_result = self.eval_manager.get_best_ptq_result()
-        if best_result and best_result.accuracy >= target_acc:
-            sess.result["target_satisfied"] = True
-            return best_result.as_dict()
-
-        if best_result is None:
-            model = fp32_model
-            applied_techniques = []
-        else:
-            if "cross_layer_equalization" not in best_result.applied_techniques:
-                sess.result["effective"] = False
-            model = best_result.load_model()
-            applied_techniques = best_result.applied_techniques
-
-        # AdaRound
-        with self.eval_manager.session("AdaRound", ptq=True) as sess:
-            model, encoding_path = self._apply_adaround(model)
-            if sess.ptq_result is None:
-                sess.set_ptq_result(model=model,
-                                    encoding_path=encoding_path,
-                                    applied_techniques=[*applied_techniques, "adaround"],
-                                    export_kwargs=self._export_kwargs)
-
-        best_result = self.eval_manager.get_best_ptq_result()
-        if best_result:
-            if "adaround" not in best_result.applied_techniques:
-                sess.result["effective"] = False
-            if best_result.accuracy >= target_acc:
-                sess.result["target_satisfied"] = True
-            return best_result.as_dict()
-
-        raise RuntimeError("None of batchnorm folding, CLE, or Adaround "
-                           "has been finished successfully.")
-
-
-@dataclass
-class PtqResult:
-    """
-    Evaluation results.
-    :param tag: Identifier string of the evaluation result.
-    :param model_path: Path to the serialized model.
-    :param encoding_path: Path to the encoding file.
-    :param accuracy: Accuracy of the model.
-    """
-    model_path: str
-    device: torch.device
-    encoding_path: str
-    accuracy: float
-    applied_techniques: List[str]
-
-    def load_model(self) -> torch.nn.Module:
-        """
-        Load model.
-        :return: Loaded model.
-        """
-        return torch.load(self.model_path).to(self.device)
-
-    def as_dict(self):
-        """Convert to dictionary"""
-        return dict(model=self.load_model(),
-                    accuracy=self.accuracy,
-                    encoding_path=self.encoding_path,
-                    applied_techniques=self.applied_techniques)
-
-
-class _EvalManager:
-    """
-    Evaluation manager for AutoQuant.
-    """
-    def __init__(self,
-                 quantsim_factory: Callable,
-                 eval_func: Callable[[torch.nn.Module], float],
-                 dummy_input_on_cpu: Union[torch.Tensor, Tuple],
-                 results_dir: str,
-                 strict_validation: bool):
-        """
-        :param quantsim_factory: A factory function that returns QuantizationSimModel.
-        :param eval_func: Evaluation function.
-        :param dummy_input: Dummy input to the model. Assumed to be located on the same device as the model.
-        :param dummy_input_on_cpu: Dummy input to the model in CPU memory.
-        :param results_dir: Base directory to save the temporary serialized model.
-        """
-        self._quantsim_factory = quantsim_factory
-        self._eval_func = eval_func
-        self._dummy_input_on_cpu = dummy_input_on_cpu
-        self._results_dir = results_dir
-        self._strict_validation = strict_validation
-
-        os.makedirs(self._results_dir, exist_ok=True)
-
-        self._all_sessions = OrderedDict() # type: OrderedDict[str, _EvalSession]
-
-    def clear(self):
-        """
-        Clear all the session status saved in the previous run
-        """
-        for sess in self._all_sessions.values():
-            sess.reset_status()
-
-    def get_best_ptq_result(self) -> Optional[PtqResult]:
-        """
-        Get the results with the highest evaluation score among the ptq results evaluated so far.
-        :return: The best evaluation result so far.
-        """
-        ptq_results = [sess.ptq_result for sess in self._all_sessions.values()
-                       if sess.ptq_result is not None]
-        if not ptq_results:
-            return None
-
-        return max(ptq_results, key=lambda ptq_result: ptq_result.accuracy)
-
-    def session(self, title: str, ptq: bool = False):
-        """
-        Session factory.
-        :param title: Title of the session.
-        :param ptq: True if this session is a ptq session
-        :return: Session object.
-        """
-        if title not in self._all_sessions:
-            session = _EvalSession(title,
-                                   self._quantsim_factory,
-                                   self._eval_func,
-                                   self._dummy_input_on_cpu,
-                                   results_dir=os.path.join(self._results_dir, ".trace"),
-                                   strict_validation=self._strict_validation,
-                                   ptq=ptq)
-            self._all_sessions[title] = session
-        return self._all_sessions[title]
-
-    HTML_TEMPLATE_FILE = os.path.join(
-        os.path.dirname(os.path.abspath(__file__)),
-        "auto_quant_diagnostics_template.html",
-    )
-
-    def export_diagnostics(self) -> str:
-        """
-        Export diagnostics in html format.
-        :return: Diagnostics string in html format.
-        """
-        loader = jinja2.FileSystemLoader(os.path.dirname(self.HTML_TEMPLATE_FILE))
-        env = jinja2.Environment(loader=loader)
-        template = env.get_template(os.path.basename(self.HTML_TEMPLATE_FILE))
-
-        if any(sess.diagnostics.contains_bokeh() for sess in self._all_sessions.values()):
-            head = CDN.render()
-        else:
-            head = ""
-
-        log = io.StringIO()
-        for sess in self._all_sessions.values():
-            if sess.diagnostics.is_empty():
-                continue
-            log.write(
-                f"<h1> {sess.title} </h1>\n"
-            )
-            content = "\n".join(
-                line.get_html_elem() for line in sess.diagnostics
-            )
-            log.write(f"{content}\n")
-
-        result = OrderedDict()
-        result["ptq_techniques"] = OrderedDict()
-
-        for sess in self._all_sessions.values():
-            if sess.is_ptq_session():
-                result["ptq_techniques"][sess.title_lowercase] = sess.result
-            else:
-                result[sess.title_lowercase] = sess.result
-
-        flowchart_metadata = _build_flowchart_metadata(result)
-
-        html = template.render(head=head, log=log.getvalue(), **flowchart_metadata)
-
-        filename = os.path.join(self._results_dir, "diagnostics.html")
-        with open(filename, "w") as f:
-            f.write(html)
-        return html
-
-
-class _EvalSession: # pylint: disable=too-many-instance-attributes
-    """
-    Evaluation session for AutoQuant.
-
-    Each session object contains a title and diagnostics produced during the session.
-    The collected diagnostics will be exported into a html file by _EvalManager.
-    """
-    def __init__(
-            self,
-            title: str,
-            quantsim_factory: Callable,
-            eval_func: Callable[[torch.nn.Module], float],
-            dummy_input_on_cpu: Union[torch.Tensor, Tuple],
-            results_dir: str,
-            strict_validation: bool,
-            ptq: bool,
-    ):
-        """
-        :param title: Title of the session.
-        :param quantsim_factory: A factory function that returns QuantizationSimModel.
-        :param eval_func: Evaluation function.
-        :param dummy_input_on_cpu: Dummy input to the model in CPU memory.
-        :param results_dir: Base directory to save the temporary serialized model.
-        :param ptq: True if this session is a ptq session
-        """
-        self.title = title
-        self._quantsim_factory = quantsim_factory
-        self._eval_func = eval_func
-        self._dummy_input_on_cpu = dummy_input_on_cpu
-        self._results_dir = results_dir
-        self._strict_validation = strict_validation
-        self._ptq = ptq
-
-        self._spinner = None
-
-        self.result = {
-            "status": None,
-            "error": None,
-            "target_satisfied": False,
-            "effective": True,
-        }
-
-        os.makedirs(self._results_dir, exist_ok=True)
-
-        self.diagnostics = Diagnostics()
-
-        # Map session title to file name.
-        # e.g. title: "Cross-Layer Equalization" -> filename: "cross_layer_equalization"
-        self.title_lowercase = self.title.lower().replace("-", " ")
-        self.title_lowercase = "_".join(self.title_lowercase.split())
-
-        stdout_write = sys.stdout.write
-        self._log = io.StringIO()
-
-        # Redirects stdout to self._log
-        def write_wrapper(*args, **kwargs):
-            self._log.write(*args, **kwargs)
-            return stdout_write(*args, **kwargs)
-
-        self._stdout_redirect = patch.object(sys.stdout, "write", write_wrapper)
-        self._ptq_result = None
-        self._cached_result = None
-
-    def is_ptq_session(self):
-        """
-        Getter method of self._ptq flag
-        """
-        return self._ptq
-
-    def reset_status(self):
-        """
-        Reset the session status saved in the previous run
-        """
-        self.result = {
-            "status": None,
-            "error": None,
-            "target_satisfied": False,
-            "effective": True,
-        }
-
-    def wrap(self, fn):
-        """
-        Return a wrapper function that caches the return value.
-
-        :param fn: Function to wrap.
-        :returns: Function whose return value is cached.
-        """
-
-        results_dir = self._results_dir
-        class CachedResult:
-            """Cached result """
-            def __init__(self, obj):
-                self._filename = os.path.join(results_dir, f".{uuid4()}")
-                while os.path.exists(self._filename):
-                    self._filename = os.path.join(results_dir, f".{uuid4()}")
-                with open(self._filename, "wb") as f:
-                    pickle.dump(obj, f)
-
-            def load(self):
-                """Load cached result """
-                with open(self._filename, "rb") as f:
-                    return pickle.load(f)
-
-        @functools.wraps(fn)
-        def wrapper(*args, **kwargs):
-            if self._cached_result:
-                return self._cached_result.load()
-            ret = fn(*args, **kwargs)
-            self._cached_result = CachedResult(ret)
-            return ret
-        return wrapper
-
-    def eval(self, model: torch.nn.Module, **kwargs):
-        """
-        Evaluate the model.
-        :param model: Model to evaluate.
-        :param **kwargs: Additional arguments to the quantsim factory.
-        :return: Eval score.
-        """
-        sim = self._quantsim_factory(model, **kwargs)
-        acc = self._eval_func(sim.model)
-        return acc
-
-    def __enter__(self):
-        self._spinner = Spinner(self.title)
-        self._spinner.__enter__()
-        self._stdout_redirect.start()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        if self._ptq_result is not None:
-            _logger.info("Session finished: %s. (eval score: %f)",
-                         self.title, self._ptq_result.accuracy)
-
-        self._spinner.__exit__(exc_type, exc_val, exc_tb)
-
-        if exc_val:
-            buffer = io.StringIO()
-            traceback.print_exception(exc_type, exc_val, exc_tb, file=buffer)
-
-            if self._strict_validation:
-                print(buffer.getvalue())
-            else:
-                print(
-                    "################################################################\n"
-                    "################################################################\n"
-                    "################################################################\n"
-                    "WARNING: The following exception was raised but ignored:\n\n"
-                    f"{buffer.getvalue()}"
-                    "################################################################\n"
-                    "################################################################\n"
-                    "################################################################\n"
-                )
-
-        self._stdout_redirect.stop()
-        self.diagnostics.add(self._log.getvalue())
-
-        self.result["error"] = exc_val
-        if not exc_val:
-            self.result["status"] = "success"
-        elif self._strict_validation:
-            self.result["status"] = "error-failed"
-        else:
-            self.result["status"] = "error-ignored"
-
-        if exc_val and not self._strict_validation:
-            # Return True so that the error doesn't propagate further
-            return True
-        return None
-
-    @property
-    def ptq_result(self) -> Optional[PtqResult]:
-        """Getter of self._ptq_result."""
-        return self._ptq_result
-
-    def set_ptq_result(
-            self,
-            applied_techniques: List[str],
-            model: torch.nn.Module = None,
-            sim: QuantizationSimModel = None,
-            acc: float = None,
-            export_kwargs: Mapping = None,
-            **kwargs
-    ) -> None:
-        """
-        Set the result of PTQ. Should be called exactly once inside a with-as block.
-
-        Exactly one among model and (sim, acc) pair should be specified.
-        1) If sim and acc is specified, save them as the result of this session.
-        2) If model is specified, evaluate the quantized accuracy of the model and save the result.
-
-        :param model: Result of PTQ.
-        :param sim: Result of PTQ. The quamtization encoding (compute_encodings()) is
-                    assumed to have been computed in advance.
-        :param acc: Eval score.
-        :param **kwargs: Additional arguments to the quantsim factory.
-        :return: None
-        """
-        if export_kwargs is None:
-            export_kwargs = {}
-
-        if sim is None:
-            assert acc is None
-            assert model is not None
-            sim = self._quantsim_factory(model, **kwargs)
-            acc = self._eval_func(sim.model)
-        else:
-            assert acc is not None
-            assert model is None
-
-        self._set_ptq_result(sim, acc, applied_techniques, export_kwargs)
-
-    def _set_ptq_result(
-            self,
-            sim: QuantizationSimModel,
-            acc: float,
-            applied_techniques: List[str],
-            export_kwargs: Mapping,
-    ) -> PtqResult:
-        """
-        Set the result of PTQ. Should be called exactly once inside a with-as block.
-
-        :param sim: Result of PTQ. The quamtization encoding (compute_encodings()) is
-                    assumed to have been computed in advance.
-        :param acc: Eval score.
-        :param export_kwargs: Additional kwargs for sim.export
-        :return: PtqResult object.
-        """
-        if self._ptq_result is not None:
-            raise RuntimeError(
-                "sess.eval() can be called only once per each _EvalSession instance."
-            )
-
-        device = utils.get_device(sim.model)
-        model_path, encoding_path = self._export(sim, export_kwargs)
-        self._ptq_result = PtqResult(
-            model_path=model_path,
-            device=device,
-            encoding_path=encoding_path,
-            accuracy=acc,
-            applied_techniques=applied_techniques,
-        )
-        return self._ptq_result
-
-    def _export(self, sim: QuantizationSimModel, export_kwargs: Mapping) -> Tuple[str, str]:
-        """
-        Export quantsim.
-        :param sim: QuantizationSimModel object to export.
-        :param export_kwargs: Additional kwargs for sim.export
-        :return: The paths where model and encoding are saved
-        """
-        sim.export(path=self._results_dir,
-                   filename_prefix=self.title_lowercase,
-                   dummy_input=self._dummy_input_on_cpu,
-                   **export_kwargs)
-        model_path = os.path.join(self._results_dir, f"{self.title_lowercase}.pth")
-        encoding_path = os.path.join(self._results_dir, f"{self.title_lowercase}.encodings")
-        _logger.info("The results of %s is saved in %s and %s.",
-                     self.title, model_path, encoding_path)
-        return model_path, encoding_path
-
-
-@contextlib.contextmanager
-def spy_auto_quant(auto_quant: AutoQuant):
-    """
-    Install a spy that collects the handles to the ptq result of
-    each stage of AutoQuant.
-
-    Typical usage::
-        >>> auto_quant = AutoQuant(...)
-        ... with auto_quant_spy(auto_quant) as spy:
-        ...     _ = auto_quant.apply(...)
-        ...
-        ... for result in spy.get_all_ptq_results():
-        ...     print(result.applied_techniques)
-        ...     print(result.accuracy)
-        ...     print(result.encoding_path)
-        ...     model = result.load_model()
-        ...     ...
-    """
-    # pylint: disable=protected-access
-    class Spy:
-        """
-        Spy that collects the handles to the ptq result of
-        each stage of AutoQuant.
-        """
-        def __init__(self, eval_manager):
-            self._eval_manager = eval_manager
-
-        def get_all_ptq_results(self) -> List[PtqResult]:
-            """Return handles to the results of AutoQuant"""
-            if self._eval_manager is None:
-                return []
-            return [sess.ptq_result for sess in self._eval_manager._all_sessions.values()
-                    if sess.ptq_result is not None]
-
-    spy = Spy(auto_quant.eval_manager)
-
-    _optimize_main = auto_quant._optimize_main
-
-    def _optimize_main_wrapper(fp32_model, target_acc):
-        return _optimize_main(fp32_model, target_acc)
-
-    try:
-        setattr(auto_quant, "_optimize_main", _optimize_main_wrapper)
-        yield spy
-    finally:
-        setattr(auto_quant, "_optimize_main", _optimize_main)
-
-
-def _build_flowchart_metadata(result: Mapping) -> Dict: # pylint: disable=too-many-return-statements
-    """
-    Build flowchart metadata for the html template of summary report
-
-    :param result: Result of AutoQuant with the following format:
-
-        result := {
-            "prepare_model": _stage_result,
-            "quantscheme_selection": _stage_result,
-            "w32_evaluation": _stage_result,
-            "ptq_techniques" [
-                "batchnorm_folding": _stage_result,
-                "cross_layer_equalization": _stage_result,
-                "adaround": _stage_result,
-            ]
-
-        }
-
-        where _stage_result is a dictionary defined as below:
-
-        _stage_result := {
-            "status": str,
-            "error": Exception,
-            "target_satisfied": bool,
-            "effective": bool,
-        }
-
-    :return: Dictionary that contains flowchart metadata for html template
-    """
-    metadata = defaultdict(str)
-    metadata.update(
-        edge_prepare_model_in='data-visited="true"',
-        node_prepare_model='data-visited="true"',
-    )
-
-    status = result['prepare_model']['status']
-    metadata.update(
-        node_prepare_model=f'data-visited="true" data-stage-result="{status}"',
-    )
-
-    if status == 'error-failed':
-        return metadata
-
-    metadata.update(
-        edge_prepare_model_out='data-visited="true"',
-    )
-
-    if "quantscheme_selection" in result:
-        status = result['quantscheme_selection']['status']
-        metadata.update(
-            node_quant_scheme_selection=f'data-visited="true" data-stage-result="{status}"',
-        )
-
-        if status == 'error-failed':
-            return metadata
-
-    metadata.update(
-        edge_quant_scheme_selection_out='data-visited="true"',
-        node_test_w32_eval_score='data-visited="true"',
-    )
-
-    if not result["w32_evaluation"]["target_satisfied"]:
-        metadata.update(
-            edge_test_w32_eval_score_if_false='data-visited="true"',
-            node_result_fail='data-visited="true"',
-        )
-        return metadata
-
-    metadata.update(
-        edge_test_w32_eval_score_if_true='data-visited="true"',
-    )
-
-
-    for ptq_name, ptq_result in result["ptq_techniques"].items():
-        status = ptq_result['status']
-        effective = ptq_result['effective']
-        if status == "success" and not effective:
-            status = "discarded"
-        metadata.update({
-            f"node_{ptq_name}": f'data-visited="true" data-stage-result="{status}"',
-        })
-
-        if status == 'error-failed':
-            return metadata
-
-        metadata.update({
-            f'edge_{ptq_name}_out': 'data-visited="true"',
-            f'node_test_{ptq_name}': 'data-visited="true"',
-        })
-
-        if ptq_result['target_satisfied']:
-            metadata.update({
-                f'edge_test_{ptq_name}_if_true': 'data-visited="true"',
-                'node_result_success': 'data-visited="true"',
-            })
-            return metadata
-
-        metadata.update({
-            f'edge_test_{ptq_name}_if_false': 'data-visited="true"',
-        })
-
-    metadata.update(
-        node_result_fail='data-visited="true"',
-    )
-
-    return metadata


Test patch:
diff --git a/TrainingExtensions/torch/test/python/test_auto_quant.py b/TrainingExtensions/torch/test/python/test_auto_quant.py
index b4a2319aaa..1d23762b16 100644
--- a/TrainingExtensions/torch/test/python/test_auto_quant.py
+++ b/TrainingExtensions/torch/test/python/test_auto_quant.py
@@ -50,7 +50,7 @@
 
 from aimet_torch import utils
 from aimet_torch.model_preparer import prepare_model
-from aimet_torch.auto_quant_v2 import AutoQuant
+from aimet_torch.auto_quant import AutoQuant
 from aimet_torch.adaround.adaround_weight import AdaroundParameters
 from aimet_torch.quantsim import QuantizationSimModel, OnnxExportApiArgs
 from aimet_torch.qc_quantize_op import StaticGridQuantWrapper
@@ -305,10 +305,10 @@ class Mocks:
         equalize_model: MagicMock
         apply_adaround: MagicMock
 
-    with patch("aimet_torch.auto_quant_v2.QuantizationSimModel", side_effect=_QuantizationSimModel) as mock_qsim,\
-            patch("aimet_torch.auto_quant_v2.fold_all_batch_norms", side_effect=bn_folding) as mock_bn_folding,\
-            patch("aimet_torch.auto_quant_v2.equalize_model", side_effect=cle) as mock_cle,\
-            patch("aimet_torch.auto_quant_v2.Adaround._apply_adaround", side_effect=adaround) as mock_adaround:
+    with patch("aimet_torch.auto_quant.QuantizationSimModel", side_effect=_QuantizationSimModel) as mock_qsim,\
+            patch("aimet_torch.auto_quant.fold_all_batch_norms", side_effect=bn_folding) as mock_bn_folding,\
+            patch("aimet_torch.auto_quant.equalize_model", side_effect=cle) as mock_cle,\
+            patch("aimet_torch.auto_quant.Adaround._apply_adaround", side_effect=adaround) as mock_adaround:
         try:
             yield Mocks(
                 eval_callback=mock_eval_callback,
@@ -384,7 +384,7 @@ def test_consecutive_calls(self, cpu_model, dummy_input, unlabeled_data_loader):
         with patch_ptq_techniques(
             bn_folded_acc, cle_acc, adaround_acc
         ) as mocks:
-            with patch("aimet_torch.auto_quant_v2.prepare_model", side_effect=prepare_model) as prepare_model_mock:
+            with patch("aimet_torch.auto_quant.prepare_model", side_effect=prepare_model) as prepare_model_mock:
                 auto_quant = AutoQuant(cpu_model,
                                        dummy_input,
                                        unlabeled_data_loader,
@@ -538,7 +538,7 @@ def error_fn(*_, **__):
                                        mocks.eval_callback,
                                        results_dir=results_dir,
                                        strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.prepare_model", side_effect=error_fn):
+                with patch("aimet_torch.auto_quant.prepare_model", side_effect=error_fn):
                     # If prepare_model fails, should return BN folding results
                     _, acc = auto_quant.run_inference()
                     assert acc == bn_folded_acc
@@ -549,7 +549,7 @@ def error_fn(*_, **__):
                                        mocks.eval_callback,
                                        results_dir=results_dir,
                                        strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.ModelValidator.validate_model", side_effect=error_fn):
+                with patch("aimet_torch.auto_quant.ModelValidator.validate_model", side_effect=error_fn):
                     # If validate_model fails, should return BN folding results
                     _, acc = auto_quant.run_inference()
                     assert acc == bn_folded_acc
@@ -560,9 +560,9 @@ def error_fn(*_, **__):
                                        mocks.eval_callback,
                                        results_dir=results_dir,
                                        strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.prepare_model", side_effect=error_fn),\
-                    patch("aimet_torch.auto_quant_v2.ModelValidator.validate_model", side_effect=error_fn),\
-                    patch("aimet_torch.auto_quant_v2.fold_all_batch_norms", side_effect=error_fn):
+                with patch("aimet_torch.auto_quant.prepare_model", side_effect=error_fn),\
+                    patch("aimet_torch.auto_quant.ModelValidator.validate_model", side_effect=error_fn),\
+                    patch("aimet_torch.auto_quant.fold_all_batch_norms", side_effect=error_fn):
                     # If all of prepare_model, validate_model, and BN folding fail, should return raw quantsim model
                     _, acc = auto_quant.run_inference()
                     assert acc == raw_quantsim_acc
@@ -587,7 +587,7 @@ def error_fn(*_, **__):
                                        mocks.eval_callback,
                                        results_dir=results_dir,
                                        strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.prepare_model", side_effect=error_fn):
+                with patch("aimet_torch.auto_quant.prepare_model", side_effect=error_fn):
                     # If prepare_model fails, should return Adaround results
                     _, acc, _ = auto_quant.optimize()
                     assert acc == adaround_acc
@@ -607,7 +607,7 @@ def error_fn(*_, **__):
                                        mocks.eval_callback,
                                        results_dir=results_dir,
                                        strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.fold_all_batch_norms", side_effect=error_fn):
+                with patch("aimet_torch.auto_quant.fold_all_batch_norms", side_effect=error_fn):
                     # If batchnorm folding fails, should return Adaround results
                     _, acc, _ = auto_quant.optimize()
                     assert acc == adaround_acc
@@ -626,7 +626,7 @@ def error_fn(*_, **__):
                                        mocks.eval_callback,
                                        results_dir=results_dir,
                                        strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.equalize_model", side_effect=error_fn):
+                with patch("aimet_torch.auto_quant.equalize_model", side_effect=error_fn):
                     # If CLE fails, should return Adaround results
                     _, acc, _ = auto_quant.optimize()
                     assert acc == adaround_acc
@@ -645,7 +645,7 @@ def error_fn(*_, **__):
                                        mocks.eval_callback,
                                        results_dir=results_dir,
                                        strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.Adaround._apply_adaround", side_effect=error_fn):
+                with patch("aimet_torch.auto_quant.Adaround._apply_adaround", side_effect=error_fn):
                     # If adaround fails, should return CLE results
                     _, acc, _ = auto_quant.optimize()
                     assert acc == cle_acc
@@ -664,9 +664,9 @@ def error_fn(*_, **__):
                                        mocks.eval_callback,
                                        results_dir=results_dir,
                                        strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.fold_all_batch_norms", side_effect=error_fn),\
-                        patch("aimet_torch.auto_quant_v2.equalize_model", side_effect=error_fn),\
-                        patch("aimet_torch.auto_quant_v2.Adaround._apply_adaround", side_effect=error_fn):
+                with patch("aimet_torch.auto_quant.fold_all_batch_norms", side_effect=error_fn),\
+                        patch("aimet_torch.auto_quant.equalize_model", side_effect=error_fn),\
+                        patch("aimet_torch.auto_quant.Adaround._apply_adaround", side_effect=error_fn):
                     # If everything fails, should raise an error
                     with pytest.raises(RuntimeError):
                         auto_quant.optimize()
@@ -685,7 +685,7 @@ def error_fn(*_, **__):
                                        mocks.eval_callback,
                                        results_dir=results_dir,
                                        strict_validation=True)
-                with patch("aimet_torch.auto_quant_v2.equalize_model", side_effect=error_fn):
+                with patch("aimet_torch.auto_quant.equalize_model", side_effect=error_fn):
                     # Hard stop if strict_validation=True
                     with pytest.raises(_Exception):
                         auto_quant.optimize()
@@ -801,7 +801,7 @@ def optimize(self, *args, **kwargs):
                 assert self._quantsim_params["quant_scheme"].output_quant_scheme == QuantScheme.post_training_tf
                 return ret
 
-            with patch("aimet_torch.auto_quant_v2.AutoQuant.optimize", optimize):
+            with patch("aimet_torch.auto_quant.AutoQuant.optimize", optimize):
                 auto_quant = AutoQuant(cpu_model,
                                        dummy_input,
                                        unlabeled_data_loader,
diff --git a/TrainingExtensions/torch/test/python/test_auto_quant_v2.py b/TrainingExtensions/torch/test/python/test_auto_quant_v2.py
deleted file mode 100644
index b4a2319aaa..0000000000
--- a/TrainingExtensions/torch/test/python/test_auto_quant_v2.py
+++ /dev/null
@@ -1,859 +0,0 @@
-# -*- mode: python -*-
-# =============================================================================
-#  @@-COPYRIGHT-START-@@
-#
-#  Copyright (c) 2022, Qualcomm Innovation Center, Inc. All rights reserved.
-#
-#  Redistribution and use in source and binary forms, with or without
-#  modification, are permitted provided that the following conditions are met:
-#
-#  1. Redistributions of source code must retain the above copyright notice,
-#     this list of conditions and the following disclaimer.
-#
-#  2. Redistributions in binary form must reproduce the above copyright notice,
-#     this list of conditions and the following disclaimer in the documentation
-#     and/or other materials provided with the distribution.
-#
-#  3. Neither the name of the copyright holder nor the names of its contributors
-#     may be used to endorse or promote products derived from this software
-#     without specific prior written permission.
-#
-#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
-#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-#  POSSIBILITY OF SUCH DAMAGE.
-#
-#  SPDX-License-Identifier: BSD-3-Clause
-#
-#  @@-COPYRIGHT-END-@@
-# =============================================================================
-
-import contextlib
-from dataclasses import dataclass
-import itertools
-from unittest.mock import patch, MagicMock
-import os
-from bs4 import BeautifulSoup
-from aimet_torch.qc_quantize_op import StaticGridQuantWrapper
-import pytest
-import shutil
-from typing import Callable
-import torch
-from torch.utils.data import Dataset, DataLoader
-
-from aimet_torch import utils
-from aimet_torch.model_preparer import prepare_model
-from aimet_torch.auto_quant_v2 import AutoQuant
-from aimet_torch.adaround.adaround_weight import AdaroundParameters
-from aimet_torch.quantsim import QuantizationSimModel, OnnxExportApiArgs
-from aimet_torch.qc_quantize_op import StaticGridQuantWrapper
-from aimet_torch.save_utils import SaveUtils
-from aimet_common.defs import QuantScheme
-
-
-class Model(torch.nn.Module):
-    """
-    Model
-    """
-
-    def __init__(self):
-        super(Model, self).__init__()
-        self._conv_0 = torch.nn.Conv2d(in_channels=3, out_channels=3, kernel_size=(3, 3), padding=1)
-        self._relu = torch.nn.ReLU()
-
-        # Test flags
-        self.register_buffer("applied_bn_folding", torch.tensor(False, dtype=torch.bool), persistent=True)
-        self.register_buffer("applied_cle", torch.tensor(False, dtype=torch.bool), persistent=True)
-        self.register_buffer("applied_adaround", torch.tensor(False, dtype=torch.bool), persistent=True)
-
-    def forward(self, x: torch.Tensor):
-        # Return the test flags along with the forward pass results so that the test flags
-        # don't get discarded when the model is converted to GraphModule by model preparer.
-        return self._relu(self._conv_0(x)),\
-               self.applied_bn_folding,\
-               self.applied_cle,\
-               self.applied_adaround
-
-
-class InvalidModel(Model):
-    def forward(self, x):
-        # This if statement throws error during model preparer
-        # since `x` is a torch.fx.Proxy object which cannot be converted to bool
-        if x[0,0,0,0]:
-            pass
-        return super().forward(x)
-
-
-@pytest.fixture(scope="session")
-def cpu_model():
-    return Model().cpu()
-
-
-@pytest.fixture(scope="session")
-def gpu_model():
-    return Model().cuda()
-
-
-@pytest.fixture(scope="session")
-def dummy_input():
-    return torch.randn((1, 3, 8, 8))
-
-
-@pytest.fixture(scope="session")
-def unlabeled_data_loader(dummy_input):
-    class MyDataset(Dataset):
-        def __init__(self, data):
-            self.data = data
-
-        def __getitem__(self, index):
-            return self.data[index]
-
-        def __len__(self):
-            return len(self.data)
-
-    dataset = MyDataset([dummy_input[0, :] for _ in range(10)])
-    return DataLoader(dataset)
-
-
-def assert_html(html_parsed, properties):
-    for id_, prop in properties.items():
-        elem = html_parsed.find(id=id_)
-        assert elem is not None
-        for prop_name, prop_val in prop.items():
-            if prop_val is None:
-                assert prop_name not in elem.attrs
-            else:
-                assert elem[prop_name] == prop_val
-
-
-_VISITED = { 'data-visited': 'true', }
-_NOT_VISITED = { 'data-visited': None, }
-_SUCCESS = {
-    'data-visited': 'true',
-    'data-stage-result': 'success'
-}
-_DISCARDED = {
-    'data-visited': 'true',
-    'data-stage-result': 'discarded'
-}
-_ERROR_IGNORED = {
-    'data-visited': 'true',
-    'data-stage-result': 'error-ignored'
-}
-_ERROR_FAILED = {
-    'data-visited': 'true',
-    'data-stage-result': 'error-failed'
-}
-
-def assert_applied_techniques(
-        output_model, acc, encoding_path,
-        target_acc, bn_folded_acc, cle_acc, adaround_acc,
-        results_dir,
-):
-    html_path = os.path.join(results_dir, 'diagnostics.html')
-    with open(html_path) as f:
-        html_parsed = BeautifulSoup(f.read(), features="html.parser")
-
-    # Batchnorm folding is always applied.
-    assert output_model.applied_bn_folding
-    assert_html(html_parsed, {
-        'node_batchnorm_folding': _SUCCESS,
-        'node_test_batchnorm_folding': _VISITED,
-    })
-
-    # If accuracy is good enough after batchnorm folding
-    if bn_folded_acc >= target_acc:
-        assert acc == bn_folded_acc
-        assert encoding_path.endswith("batchnorm_folding.encodings")
-        assert not output_model.applied_cle
-        assert not output_model.applied_adaround
-
-        assert_html(html_parsed, {
-            'node_cross_layer_equalization': _NOT_VISITED,
-            'node_test_cross_layer_equalization': _NOT_VISITED,
-            'node_adaround': _NOT_VISITED,
-            'node_test_adaround': _NOT_VISITED,
-            'node_result_fail': _NOT_VISITED,
-            'node_result_success': _VISITED,
-        })
-        return
-
-    # CLE should be applied if and only if it brings accuracy gain
-    assert output_model.applied_cle == (bn_folded_acc < cle_acc)
-
-    assert_html(html_parsed, {
-        'node_cross_layer_equalization': _SUCCESS if output_model.applied_cle else _DISCARDED,
-        'node_test_cross_layer_equalization': _VISITED,
-    })
-
-    # If accuracy is good enough after cle
-    if cle_acc >= target_acc:
-        assert acc == cle_acc
-        assert encoding_path.endswith("cross_layer_equalization.encodings")
-        assert output_model.applied_cle
-        assert not output_model.applied_adaround
-
-        assert_html(html_parsed, {
-            'node_adaround': _NOT_VISITED,
-            'node_test_adaround': _NOT_VISITED,
-            'node_result_fail': _NOT_VISITED,
-            'node_result_success': _VISITED,
-        })
-        return
-
-    assert output_model.applied_adaround == (adaround_acc >= max(bn_folded_acc, cle_acc))
-
-    assert_html(html_parsed, {
-        'node_adaround': _SUCCESS if output_model.applied_adaround else _DISCARDED,
-        'node_test_adaround': _VISITED,
-    })
-
-    # If accuracy is good enough after adaround
-    if adaround_acc >= target_acc:
-        assert acc == adaround_acc
-        assert encoding_path.endswith("adaround.encodings")
-        assert output_model.applied_adaround
-
-        assert_html(html_parsed, {
-            'node_result_fail': _NOT_VISITED,
-            'node_result_success': _VISITED,
-        })
-        return
-
-    assert_html(html_parsed, {
-        'node_result_fail': _VISITED,
-        'node_result_success': _NOT_VISITED,
-    })
-
-    assert acc == max(bn_folded_acc, cle_acc, adaround_acc)
-
-    if max(bn_folded_acc, cle_acc, adaround_acc) == bn_folded_acc:
-        assert encoding_path.endswith("batchnorm_folding.encodings")
-    elif max(bn_folded_acc, cle_acc, adaround_acc) == cle_acc:
-        assert encoding_path.endswith("cross_layer_equalization.encodings")
-    else:
-        assert encoding_path.endswith("adaround.encodings")
-
-
-FP32_ACC = .8
-W32_ACC = FP32_ACC # Assume W32 accuracy is equal to FP32 accuracy
-RAW_QUANTSIM_ACC = 0.1
-
-
-@contextlib.contextmanager
-def patch_ptq_techniques(bn_folded_acc, cle_acc, adaround_acc, fp32_acc=None, w32_acc=None, raw_quantsim_acc=None):
-    if fp32_acc is None:
-        fp32_acc = FP32_ACC
-
-    if w32_acc is None:
-        w32_acc = W32_ACC
-
-    if raw_quantsim_acc is None:
-        raw_quantsim_acc = RAW_QUANTSIM_ACC
-
-    const_true = torch.tensor(True, dtype=torch.bool)
-
-    def bn_folding(model: Model, *_, **__):
-        model.applied_bn_folding.copy_(const_true)
-        return tuple()
-
-    def cle(model: Model, *_, **__):
-        model.applied_bn_folding.copy_(const_true)
-        model.applied_cle.copy_(const_true)
-
-    def adaround(sim, *_, **__):
-        sim.model.applied_adaround.copy_(const_true)
-        SaveUtils.remove_quantization_wrappers(sim.model)
-        return sim.model
-
-    class _QuantizationSimModel(QuantizationSimModel):
-        def compute_encodings(self, *_):
-            pass
-
-        def set_and_freeze_param_encodings(self, _):
-            pass
-
-    def mock_eval_callback(model, _):
-        if not isinstance(model._conv_0, StaticGridQuantWrapper):
-            # Not quantized: return fp32 accuracy
-            return fp32_acc
-        if model._conv_0.param_quantizers["weight"].bitwidth == 32:
-            # W32 evaluation for early exit. Return W32 accuracy
-            return w32_acc
-
-        acc = raw_quantsim_acc
-        if model.applied_bn_folding:
-            acc = bn_folded_acc
-        if model.applied_cle:
-            acc = cle_acc
-        if model.applied_adaround:
-            acc = adaround_acc
-        return acc
-
-    @dataclass
-    class Mocks:
-        eval_callback: Callable
-        QuantizationSimModel: MagicMock
-        fold_all_batch_norms: MagicMock
-        equalize_model: MagicMock
-        apply_adaround: MagicMock
-
-    with patch("aimet_torch.auto_quant_v2.QuantizationSimModel", side_effect=_QuantizationSimModel) as mock_qsim,\
-            patch("aimet_torch.auto_quant_v2.fold_all_batch_norms", side_effect=bn_folding) as mock_bn_folding,\
-            patch("aimet_torch.auto_quant_v2.equalize_model", side_effect=cle) as mock_cle,\
-            patch("aimet_torch.auto_quant_v2.Adaround._apply_adaround", side_effect=adaround) as mock_adaround:
-        try:
-            yield Mocks(
-                eval_callback=mock_eval_callback,
-                QuantizationSimModel=mock_qsim,
-                fold_all_batch_norms=mock_bn_folding,
-                equalize_model=mock_cle,
-                apply_adaround=mock_adaround,
-            )
-        finally:
-            pass
-
-
-class TestAutoQuant:
-    def test_auto_quant_run_inference(self, cpu_model, dummy_input, unlabeled_data_loader):
-        bn_folded_acc = .5
-
-        with patch_ptq_techniques(
-            bn_folded_acc, None, None
-        ) as mocks:
-            with create_tmp_directory() as results_dir:
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir)
-                auto_quant.run_inference()
-
-    @pytest.mark.parametrize(
-        "bn_folded_acc, cle_acc, adaround_acc",
-        itertools.permutations([.5, .6, .7])
-    )
-    @pytest.mark.parametrize("allowed_accuracy_drop", [.05, .15])
-    def test_auto_quant_cpu(
-            self, cpu_model, dummy_input, unlabeled_data_loader,
-            allowed_accuracy_drop, bn_folded_acc, cle_acc, adaround_acc,
-    ):
-        self._test_auto_quant(
-            cpu_model, dummy_input, unlabeled_data_loader,
-            allowed_accuracy_drop, bn_folded_acc, cle_acc, adaround_acc,
-        )
-
-    @pytest.mark.cuda
-    def test_auto_quant_gpu(self, gpu_model, dummy_input, unlabeled_data_loader):
-        bn_folded_acc, cle_acc, adaround_acc = .5, .6, .7
-        allowed_accuracy_drop = .15
-
-        self._test_auto_quant(
-            gpu_model, dummy_input.cuda(), unlabeled_data_loader,
-            allowed_accuracy_drop, bn_folded_acc, cle_acc, adaround_acc,
-        )
-
-    def test_consecutive_calls(self, cpu_model, dummy_input, unlabeled_data_loader):
-        bn_folded_acc, cle_acc, adaround_acc = .5, .6, .7
-
-        with patch_ptq_techniques(
-            bn_folded_acc, cle_acc, adaround_acc
-        ) as mocks:
-            with create_tmp_directory() as results_dir:
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir)
-
-                # Should return proper model & summary report
-                # regardless of consecutive calls
-                for allowed_accuracy_drop in (.5, .4, .3, .2, .1, .05):
-                    self._do_test_optimize_auto_quant(
-                        auto_quant, cpu_model,
-                        allowed_accuracy_drop, bn_folded_acc, cle_acc, adaround_acc
-                    )
-
-        with patch_ptq_techniques(
-            bn_folded_acc, cle_acc, adaround_acc
-        ) as mocks:
-            with patch("aimet_torch.auto_quant_v2.prepare_model", side_effect=prepare_model) as prepare_model_mock:
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir)
-
-                # When run_inference() and optimize() are called in back-to-back,
-                # reusable intermediate reseults should be always reused.
-                auto_quant.run_inference()
-                auto_quant.optimize()
-                assert prepare_model_mock.call_count == 1
-                assert mocks.fold_all_batch_norms.call_count == 1
-                assert mocks.equalize_model.call_count == 1
-
-                auto_quant.optimize()
-                assert prepare_model_mock.call_count == 1
-                assert mocks.fold_all_batch_norms.call_count == 1
-                assert mocks.equalize_model.call_count == 1
-
-                self._do_test_optimize_auto_quant(
-                    auto_quant, cpu_model,
-                    0.0, bn_folded_acc, cle_acc, adaround_acc
-                )
-                assert prepare_model_mock.call_count == 1
-                assert mocks.fold_all_batch_norms.call_count == 1
-                assert mocks.equalize_model.call_count == 1
-
-    def _test_auto_quant(
-            self, model, dummy_input, unlabeled_data_loader,
-            allowed_accuracy_drop, bn_folded_acc, cle_acc, adaround_acc,
-    ):
-        with patch_ptq_techniques(
-            bn_folded_acc, cle_acc, adaround_acc
-        ) as mocks:
-            with create_tmp_directory() as results_dir:
-                auto_quant = AutoQuant(model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir)
-                self._do_test_optimize_auto_quant(
-                    auto_quant, model, allowed_accuracy_drop,
-                    bn_folded_acc, cle_acc, adaround_acc
-                )
-
-    def _do_test_optimize_auto_quant(
-            self, auto_quant, input_model,
-            allowed_accuracy_drop, bn_folded_acc, cle_acc, adaround_acc,
-    ):
-            target_acc = FP32_ACC - allowed_accuracy_drop
-
-            output_model, acc, encoding_path = auto_quant.optimize(allowed_accuracy_drop)
-
-            assert utils.get_device(output_model) == utils.get_device(input_model)
-            assert_applied_techniques(
-                output_model, acc, encoding_path,
-                target_acc, bn_folded_acc, cle_acc, adaround_acc,
-                auto_quant.results_dir,
-            )
-
-    def test_auto_quant_invalid_input(self, cpu_model, dummy_input, unlabeled_data_loader):
-        with pytest.raises(ValueError):
-            AutoQuant(None, dummy_input, unlabeled_data_loader, lambda: None)
-
-        with pytest.raises(ValueError):
-            AutoQuant(cpu_model, None, unlabeled_data_loader, lambda: None)
-
-        with pytest.raises(ValueError):
-            AutoQuant(cpu_model, dummy_input, None, lambda: None)
-
-        with pytest.raises(ValueError):
-            AutoQuant(cpu_model, dummy_input, unlabeled_data_loader, None)
-
-        with pytest.raises(ValueError):
-            AutoQuant(cpu_model, dummy_input, unlabeled_data_loader, lambda: None, results_dir=None)
-
-        with pytest.raises(ValueError):
-            AutoQuant(cpu_model, dummy_input, unlabeled_data_loader, lambda: None, strict_validation=None)
-
-        # Bitwidth < 4 or bitwidth > 32
-        with pytest.raises(ValueError):
-            AutoQuant(cpu_model, dummy_input, unlabeled_data_loader, lambda: None, param_bw=2)
-
-        with pytest.raises(ValueError):
-            AutoQuant(cpu_model, dummy_input, unlabeled_data_loader, lambda: None, param_bw=64)
-
-        with pytest.raises(ValueError):
-            AutoQuant(cpu_model, dummy_input, unlabeled_data_loader, lambda: None, output_bw=2)
-
-        with pytest.raises(ValueError):
-            AutoQuant(cpu_model, dummy_input, unlabeled_data_loader, lambda: None, output_bw=64)
-
-        auto_quant = AutoQuant(cpu_model, dummy_input, unlabeled_data_loader, lambda: None)
-        # Allowed accuracy drop < 0
-        with pytest.raises(ValueError):
-            _ = auto_quant.optimize(-1.0)
-
-    def test_auto_quant_model_preparer(self, unlabeled_data_loader, dummy_input):
-        bn_folded_acc, cle_acc, adaround_acc = .4, .5, .6
-
-        with patch_ptq_techniques(
-            bn_folded_acc, cle_acc, adaround_acc
-        ) as mocks:
-            auto_quant = AutoQuant(InvalidModel(),
-                                   dummy_input,
-                                   unlabeled_data_loader,
-                                   mocks.eval_callback,
-                                   strict_validation=True)
-
-            # If strict_validation is True (default), AutoQuant crashes with an exception.
-            with pytest.raises(torch.fx.proxy.TraceError):
-                auto_quant.run_inference()
-
-            # If strict_validation is True (default), AutoQuant crashes with an exception.
-            with pytest.raises(torch.fx.proxy.TraceError):
-                auto_quant.optimize()
-
-        with patch_ptq_techniques(
-            bn_folded_acc, cle_acc, adaround_acc
-        ) as mocks:
-            auto_quant = AutoQuant(InvalidModel(),
-                                   dummy_input,
-                                   unlabeled_data_loader,
-                                   mocks.eval_callback,
-                                   strict_validation=False)
-
-            # If strict_validation is False, AutoQuant ignores the errors and proceed. 
-            auto_quant.run_inference()
-
-            # If strict_validation is False, AutoQuant ignores the errors and proceed. 
-            auto_quant.optimize()
-
-    def test_auto_quant_inference_fallback(
-        self, cpu_model, dummy_input, unlabeled_data_loader,
-    ):
-        class _Exception(Exception):
-            pass
-
-        def error_fn(*_, **__):
-            raise _Exception
-
-        bn_folded_acc = .4
-        raw_quantsim_acc = bn_folded_acc + 1e-5
-        with patch_ptq_techniques(
-            bn_folded_acc, None, None, raw_quantsim_acc=raw_quantsim_acc
-        ) as mocks:
-            with create_tmp_directory() as results_dir:
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir,
-                                       strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.prepare_model", side_effect=error_fn):
-                    # If prepare_model fails, should return BN folding results
-                    _, acc = auto_quant.run_inference()
-                    assert acc == bn_folded_acc
-
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir,
-                                       strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.ModelValidator.validate_model", side_effect=error_fn):
-                    # If validate_model fails, should return BN folding results
-                    _, acc = auto_quant.run_inference()
-                    assert acc == bn_folded_acc
-
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir,
-                                       strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.prepare_model", side_effect=error_fn),\
-                    patch("aimet_torch.auto_quant_v2.ModelValidator.validate_model", side_effect=error_fn),\
-                    patch("aimet_torch.auto_quant_v2.fold_all_batch_norms", side_effect=error_fn):
-                    # If all of prepare_model, validate_model, and BN folding fail, should return raw quantsim model
-                    _, acc = auto_quant.run_inference()
-                    assert acc == raw_quantsim_acc
-
-    def test_auto_quant_optimize_fallback(
-        self, cpu_model, dummy_input, unlabeled_data_loader,
-    ):
-        class _Exception(Exception):
-            pass
-
-        def error_fn(*_, **__):
-            raise _Exception
-
-        bn_folded_acc, cle_acc, adaround_acc = .4, .5, .6
-        with patch_ptq_techniques(
-            bn_folded_acc, cle_acc, adaround_acc
-        ) as mocks:
-            with create_tmp_directory() as results_dir:
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir,
-                                       strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.prepare_model", side_effect=error_fn):
-                    # If prepare_model fails, should return Adaround results
-                    _, acc, _ = auto_quant.optimize()
-                    assert acc == adaround_acc
-
-                    with open(os.path.join(results_dir, 'diagnostics.html')) as f:
-                        html_parsed = BeautifulSoup(f.read(), features="html.parser")
-                        assert_html(html_parsed, {
-                            'node_prepare_model': _ERROR_IGNORED,
-                            'node_batchnorm_folding': _SUCCESS,
-                            'node_cross_layer_equalization': _SUCCESS,
-                            'node_adaround': _SUCCESS,
-                        })
-
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir,
-                                       strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.fold_all_batch_norms", side_effect=error_fn):
-                    # If batchnorm folding fails, should return Adaround results
-                    _, acc, _ = auto_quant.optimize()
-                    assert acc == adaround_acc
-
-                    with open(os.path.join(results_dir, 'diagnostics.html')) as f:
-                        html_parsed = BeautifulSoup(f.read(), features="html.parser")
-                        assert_html(html_parsed, {
-                            'node_batchnorm_folding': _ERROR_IGNORED,
-                            'node_cross_layer_equalization': _SUCCESS,
-                            'node_adaround': _SUCCESS,
-                        })
-
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir,
-                                       strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.equalize_model", side_effect=error_fn):
-                    # If CLE fails, should return Adaround results
-                    _, acc, _ = auto_quant.optimize()
-                    assert acc == adaround_acc
-
-                    with open(os.path.join(results_dir, 'diagnostics.html')) as f:
-                        html_parsed = BeautifulSoup(f.read(), features="html.parser")
-                        assert_html(html_parsed, {
-                            'node_batchnorm_folding': _SUCCESS,
-                            'node_cross_layer_equalization': _ERROR_IGNORED,
-                            'node_adaround': _SUCCESS,
-                        })
-
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir,
-                                       strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.Adaround._apply_adaround", side_effect=error_fn):
-                    # If adaround fails, should return CLE results
-                    _, acc, _ = auto_quant.optimize()
-                    assert acc == cle_acc
-
-                    with open(os.path.join(results_dir, 'diagnostics.html')) as f:
-                        html_parsed = BeautifulSoup(f.read(), features="html.parser")
-                        assert_html(html_parsed, {
-                            'node_batchnorm_folding': _SUCCESS,
-                            'node_cross_layer_equalization': _SUCCESS,
-                            'node_adaround': _ERROR_IGNORED,
-                        })
-
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir,
-                                       strict_validation=False)
-                with patch("aimet_torch.auto_quant_v2.fold_all_batch_norms", side_effect=error_fn),\
-                        patch("aimet_torch.auto_quant_v2.equalize_model", side_effect=error_fn),\
-                        patch("aimet_torch.auto_quant_v2.Adaround._apply_adaround", side_effect=error_fn):
-                    # If everything fails, should raise an error
-                    with pytest.raises(RuntimeError):
-                        auto_quant.optimize()
-
-                    with open(os.path.join(results_dir, 'diagnostics.html')) as f:
-                        html_parsed = BeautifulSoup(f.read(), features="html.parser")
-                        assert_html(html_parsed, {
-                            'node_batchnorm_folding': _ERROR_IGNORED,
-                            'node_cross_layer_equalization': _ERROR_IGNORED,
-                            'node_adaround': _ERROR_IGNORED,
-                        })
-
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir,
-                                       strict_validation=True)
-                with patch("aimet_torch.auto_quant_v2.equalize_model", side_effect=error_fn):
-                    # Hard stop if strict_validation=True
-                    with pytest.raises(_Exception):
-                        auto_quant.optimize()
-
-                    with open(os.path.join(results_dir, 'diagnostics.html')) as f:
-                        html_parsed = BeautifulSoup(f.read(), features="html.parser")
-                        assert_html(html_parsed, {
-                            'node_batchnorm_folding': _SUCCESS,
-                            'node_cross_layer_equalization': _ERROR_FAILED,
-                            'node_adaround': _NOT_VISITED,
-                        })
-
-    def test_auto_quant_early_exit(self, cpu_model, dummy_input, unlabeled_data_loader):
-        allowed_accuracy_drop = 0.1
-        w32_acc = FP32_ACC - (allowed_accuracy_drop * 2)
-
-        with create_tmp_directory() as results_dir:
-            with patch_ptq_techniques(
-                bn_folded_acc=0, cle_acc=0, adaround_acc=0, w32_acc=w32_acc
-            ) as mocks:
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir)
-                output_model, acc, encoding_path = auto_quant.optimize(allowed_accuracy_drop)
-
-            assert output_model is None
-            assert acc is None
-            assert encoding_path is None
-
-            with open(os.path.join(results_dir, 'diagnostics.html')) as f:
-                html_parsed = BeautifulSoup(f.read(), features="html.parser")
-                assert_html(html_parsed, {
-                    'node_test_w32_eval_score': _VISITED,
-                    'node_batchnorm_folding': _NOT_VISITED,
-                    'node_cross_layer_equalization': _NOT_VISITED,
-                    'node_adaround': _NOT_VISITED,
-                    'node_result_fail': _VISITED,
-                })
-
-    def test_auto_quant_caching(
-        self, cpu_model, dummy_input, unlabeled_data_loader,
-    ):
-        allowed_accuracy_drop = 0.0
-        bn_folded_acc, cle_acc, adaround_acc = .4, .5, .6
-        cache_id = "unittest"
-
-        with patch_ptq_techniques(
-            bn_folded_acc, cle_acc, adaround_acc
-        ) as mocks:
-            with create_tmp_directory() as results_dir:
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir,
-                                       cache_id=cache_id)
-
-                cache_files  = [
-                    os.path.join(results_dir, ".auto_quant_cache", cache_id, f"{key}.pkl")
-                    for key in ("batchnorm_folding", "cle", "adaround")
-                ]
-
-                # No previously cached results
-                auto_quant.optimize(allowed_accuracy_drop)
-
-                for cache_file in cache_files:
-                    assert os.path.exists(cache_file)
-
-                assert mocks.fold_all_batch_norms.call_count == 1
-                assert mocks.equalize_model.call_count == 1
-                assert mocks.apply_adaround.call_count == 1
-
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback,
-                                       results_dir=results_dir,
-                                       cache_id=cache_id)
-                # Load cached result
-                auto_quant.optimize(allowed_accuracy_drop)
-
-                # PTQ functions should not be called twice.
-                assert mocks.fold_all_batch_norms.call_count == 1
-                assert mocks.equalize_model.call_count == 1
-                assert mocks.apply_adaround.call_count == 1
-
-    def test_auto_quant_scheme_selection(
-        self, cpu_model, dummy_input, unlabeled_data_loader,
-    ):
-        allowed_accuracy_drop = 0.0
-        bn_folded_acc, cle_acc, adaround_acc = .4, .5, .6
-        with patch_ptq_techniques(
-            bn_folded_acc, cle_acc, adaround_acc
-        ) as mocks:
-            def eval_callback(model, _):
-                # Assumes the model's eval score drops to zero
-                # unless param_quant_scheme == tfe and output_quant_scheme == tf
-                if isinstance(model._conv_0, StaticGridQuantWrapper):
-                    if model._conv_0.param_quantizers["weight"].quant_scheme != QuantScheme.post_training_tf_enhanced:
-                        return 0.0
-                    if model._conv_0.output_quantizers[0].quant_scheme != QuantScheme.post_training_tf:
-                        return 0.0
-                return mocks.eval_callback(model, _)
-
-            _optimize = AutoQuant.optimize
-            def optimize(self, *args, **kwargs):
-                # Since all the other candidates (tf-tf, tfe-tfe, and tfe-percentile) yields zero accuracy,
-                # it is expected that tf-tfe is selected as the quant scheme for AutoQuant.
-                ret = _optimize(self, *args, **kwargs)
-                assert self._quantsim_params["quant_scheme"].param_quant_scheme == QuantScheme.post_training_tf_enhanced
-                assert self._quantsim_params["quant_scheme"].output_quant_scheme == QuantScheme.post_training_tf
-                return ret
-
-            with patch("aimet_torch.auto_quant_v2.AutoQuant.optimize", optimize):
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       eval_callback)
-                auto_quant.optimize(allowed_accuracy_drop)
-
-    def test_set_additional_params(self, cpu_model, dummy_input, unlabeled_data_loader):
-        allowed_accuracy_drop = 0
-        bn_folded_acc = .1
-        cle_acc = .2
-        adaround_acc = .3
-        with patch_ptq_techniques(bn_folded_acc, cle_acc, adaround_acc) as mocks:
-            export = QuantizationSimModel.export
-
-            def export_wrapper(*args, **kwargs):
-                assert kwargs["onnx_export_args"].opset_version == 10
-                assert kwargs["propagate_encodings"]
-                return export(*args, **kwargs)
-
-            try:
-                setattr(QuantizationSimModel, "export", export_wrapper)
-                auto_quant = AutoQuant(cpu_model,
-                                       dummy_input,
-                                       unlabeled_data_loader,
-                                       mocks.eval_callback)
-                adaround_params = AdaroundParameters(unlabeled_data_loader, 1)
-                auto_quant.set_adaround_params(adaround_params)
-
-                auto_quant.set_export_params(OnnxExportApiArgs(10), True)
-
-                self._do_test_optimize_auto_quant(
-                    auto_quant, cpu_model,
-                    allowed_accuracy_drop, bn_folded_acc, cle_acc, adaround_acc
-                )
-                adaround_args, _ = mocks.apply_adaround.call_args
-                _, _, _, actual_adaround_params = adaround_args
-                assert adaround_params == actual_adaround_params
-            finally:
-                setattr(QuantizationSimModel, "export", export)
-
-
-@contextlib.contextmanager
-def create_tmp_directory(dirname: str = "/tmp/.aimet_unittest"):
-    success = False
-    try:
-        os.makedirs(dirname, exist_ok=True)
-        success = True
-    except FileExistsError:
-        raise
-
-    try:
-        yield dirname
-    finally:
-        if success:
-            shutil.rmtree(dirname)