fix format

Signed-off-by: binliu <[email protected]>

Author: binliunls

monai/handlers/__init__.py

@@ -29,8 +29,8 @@
 from .metrics_reloaded_handler import MetricsReloadedBinaryHandler, MetricsReloadedCategoricalHandler
 from .metrics_saver import MetricsSaver
 from .mlflow_handler import MLFlowHandler
-from .model_quantizer import ModelQuantizer
 from .model_calibrator import ModelCalibrater
+from .model_quantizer import ModelQuantizer
 from .nvtx_handlers import MarkHandler, RangeHandler, RangePopHandler, RangePushHandler
 from .panoptic_quality import PanopticQuality
 from .parameter_scheduler import ParamSchedulerHandler

monai/handlers/model_calibrator.py

@@ -11,13 +11,13 @@
 
 from __future__ import annotations
 
+from functools import partial
 from typing import TYPE_CHECKING
 
-import torch
 import modelopt.torch.quantization as mtq
-from functools import partial
-from monai.utils import IgniteInfo, min_version, optional_import
+import torch
 
+from monai.utils import IgniteInfo, min_version, optional_import
 
 Events, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Events")
 Checkpoint, _ = optional_import("ignite.handlers", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Checkpoint")
@@ -40,13 +40,7 @@ class ModelCalibrater:
 
     """
 
-    def __init__(
-        self,
-        model: torch.nn.Module,
-        export_path: str,
-        config: dict= mtq.INT8_SMOOTHQUANT_CFG,
-
-    ) -> None:
+    def __init__(self, model: torch.nn.Module, export_path: str, config: dict = mtq.INT8_SMOOTHQUANT_CFG) -> None:
         self.model = model
         self.export_path = export_path
         self.config = config
@@ -65,4 +59,4 @@ def _model_wrapper(engine, model):
     def __call__(self, engine) -> None:
         quant_fun = partial(self._model_wrapper, engine)
         model = mtq.quantize(self.model, self.config, quant_fun)
-        torch.save(self.model.state_dict(), self.export_path)
+        torch.save(model.state_dict(), self.export_path)

monai/handlers/model_quantizer.py

@@ -11,21 +11,16 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from types import MethodType
 from typing import TYPE_CHECKING
-from collections.abc import Sequence
 
 import torch
+from torch.ao.quantization.quantize_pt2e import prepare_qat_pt2e
+from torch.ao.quantization.quantizer import Quantizer
+from torch.ao.quantization.quantizer.xnnpack_quantizer import XNNPACKQuantizer, get_symmetric_quantization_config
 
 from monai.utils import IgniteInfo, min_version, optional_import
-from torch.ao.quantization.quantizer import Quantizer
-from torch.ao.quantization.quantizer.xnnpack_quantizer import (
-  XNNPACKQuantizer,
-  get_symmetric_quantization_config,
-)
-from torch.ao.quantization.quantize_pt2e import (
-  prepare_qat_pt2e,
-)
 
 Events, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Events")
 Checkpoint, _ = optional_import("ignite.handlers", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Checkpoint")
@@ -49,17 +44,14 @@ class ModelQuantizer:
     """
 
     def __init__(
-        self,
-        model: torch.nn.Module,
-        example_inputs: Sequence,
-        export_path: str,
-        quantizer: Quantizer | None = None,
-
+        self, model: torch.nn.Module, example_inputs: Sequence, export_path: str, quantizer: Quantizer | None = None
     ) -> None:
         self.model = model
         self.example_inputs = example_inputs
         self.export_path = export_path
-        self.quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config()) if quantizer is None else quantizer
+        self.quantizer = (
+            XNNPACKQuantizer().set_global(get_symmetric_quantization_config()) if quantizer is None else quantizer
+        )
 
     def attach(self, engine: Engine) -> None:
         """