fix:style corrections

quanto/quantize.py

@@ -12,9 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .nn import QModuleMixin, quantize_module
-import torch
 from typing import Dict, Union
+
+import torch
+
+from .nn import QModuleMixin, quantize_module
+
+
 __all__ = ["quantize", "freeze"]
 
 
@@ -56,4 +60,4 @@ def freeze(model):
 
 def requantize(model: torch.nn.Module, state_dict: Dict[str, Union[torch.Tensor, str]]):
     quantize(model)  # quantize the model from tensor
-    model.load_state_dict(state_dict, assign=True)  #reload the state_dict which is already quantized
+    model.load_state_dict(state_dict, assign=True)  # reload the state_dict which is already quantized

test/model/test_requantize_mlp.py

@@ -3,18 +3,13 @@
 
 import pytest
 import torch
-from helpers import assert_similar, get_device_memory, random_qactivation, random_tensor
+from helpers import random_qactivation, random_tensor
 
 from quanto import (
-    AbsmaxOptimizer,
     Calibration,
-    MaxOptimizer,
     QLinear,
     QTensor,
     freeze,
-    qfloat8_e4m3fn,
-    qfloat8_e5m2,
-    qint4,
     qint8,
     quantize,
     safe_load,
@@ -23,6 +18,7 @@
 from quanto.nn import QModuleMixin
 from quanto.quantize import requantize
 
+
 class MLP(torch.nn.Module):
     def __init__(self, input_size, output_size, hidden_size):
         super().__init__()
@@ -82,16 +78,11 @@ def test_serialize_quantized_mlp(weights, dtype, serialization, device):
     freeze(model)
     state_dict = save_and_reload_state_dict(model.state_dict(), serialization)
     model_reloaded = MLP(input_features, hidden_features, output_features).to(device)
-    requantize(model_reloaded,state_dict) 
+    requantize(model_reloaded, state_dict)
     for name, module in model.named_modules():
         if isinstance(module, QModuleMixin):
             module_reloaded = getattr(model_reloaded, name)
             assert torch.equal(module_reloaded.weight._data, module.weight._data)
             assert torch.equal(module_reloaded.weight._scale, module.weight._scale)
             assert torch.equal(module_reloaded.input_scale, module.input_scale)
             assert torch.equal(module_reloaded.output_scale, module.output_scale)
-
-
-
-
-