ruff

Author: jerryzh168

torchao/dtypes/uintx/gemlite_layout.py

@@ -1,3 +1,4 @@
+import warnings
 from dataclasses import dataclass
 from typing import Dict, Optional, Tuple
 
@@ -15,7 +16,6 @@
 from torchao.dtypes.utils import Layout, is_device
 from torchao.quantization.quant_primitives import quantize_affine
 from torchao.utils import fill_defaults
-import warnings
 
 aten = torch.ops.aten
 
@@ -81,7 +81,8 @@ def apply_gemlite_quant(
         warnings.simplefilter("once", UserWarning)
         warnings.warn(
             "Gemlite only works for layers with in_features or out_features divisible by 128, "
-            + "some layers have been skipped", UserWarning
+            + "some layers have been skipped",
+            UserWarning,
         )
         return weight
 
@@ -182,9 +183,10 @@ def from_plain(
             exhaustive=False,
             use_cuda_graph=False,
         )
-        if _layout.group_size == None and _layout.bit_width == 4:
+        if _layout.group_size is None and _layout.bit_width == 4:
             from gemlite.core import GEMLITE_ACC_DTYPE
             from gemlite.dtypes import DType
+
             GEMLITE_ACC_DTYPE[DType.FP16] = DType.FP32
 
         out_features, in_features = int_data.shape