lint

Author: petrex

test/float8/test_base.py

@@ -25,7 +25,6 @@
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
 
-
 from torchao.float8.config import (
     CastConfig,
     Float8LinearConfig,

test/integration/test_integration.py

@@ -80,8 +80,8 @@
     benchmark_model,
     is_fbcode,
     is_sm_at_least_90,
-    unwrap_tensor_subclass,
     skip_if_rocm,
+    unwrap_tensor_subclass,
 )
 
 try:

test/kernel/test_galore_downproj.py

@@ -12,6 +12,7 @@
 from torchao.prototype.galore.kernels.matmul import set_tuner_top_k as matmul_tuner_topk
 from torchao.prototype.galore.kernels.matmul import triton_mm_launcher
 from torchao.utils import skip_if_rocm
+
 torch.manual_seed(0)
 
 matmul_tuner_topk(10)

test/prototype/test_awq.py

@@ -5,13 +5,16 @@
 import torch
 
 from torchao.quantization import quantize_
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_3, TORCH_VERSION_AT_LEAST_2_5, skip_if_rocm
+from torchao.utils import (
+    TORCH_VERSION_AT_LEAST_2_3,
+    TORCH_VERSION_AT_LEAST_2_5,
+    skip_if_rocm,
+)
 
 if TORCH_VERSION_AT_LEAST_2_3:
     from torchao.prototype.awq import AWQObservedLinear, awq_uintx, insert_awq_observer_
 
 
-
 class ToyLinearModel(torch.nn.Module):
     def __init__(self, m=512, n=256, k=128):
         super().__init__()

test/quantization/test_galore_quant.py

@@ -19,6 +19,7 @@
     triton_quantize_blockwise,
 )
 from torchao.utils import skip_if_rocm
+
 SEED = 0
 torch.manual_seed(SEED)
 

test/test_utils.py

@@ -6,9 +6,6 @@
 from torchao.utils import TorchAOBaseTensor, torch_version_at_least
 
 
-
-
-
 class TestTorchVersionAtLeast(unittest.TestCase):
     def test_torch_version_at_least(self):
         test_cases = [

torchao/utils.py

@@ -7,6 +7,7 @@
 from math import gcd
 from typing import Any, Callable, Tuple
 
+import pytest
 import torch
 import torch.nn.utils.parametrize as parametrize
 
@@ -160,6 +161,7 @@ def wrapper(*args, **kwargs):
 
     return decorator
 
+
 def skip_if_rocm(message=None):
     """Decorator to skip tests on ROCm platform with custom message.
 
@@ -186,6 +188,7 @@ def wrapper(*args, **kwargs):
         return decorator(func)
     return decorator
 
+
 def compute_max_diff(output: torch.Tensor, output_ref: torch.Tensor) -> torch.Tensor:
     return torch.mean(torch.abs(output - output_ref)) / torch.mean(
         torch.abs(output_ref)