Review comments from @kylesayrs compare enum directly

Signed-off-by: Rahul Tuli <rahul@neuralmagic.com>
vllm-project · Dec 3, 2024 · afc0b5f · afc0b5f
1 parent 0bd66c8
commit afc0b5f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 6 deletions.
diff --git a/src/llmcompressor/transformers/compression/quantization_format.py b/src/llmcompressor/transformers/compression/quantization_format.py
@@ -37,8 +37,7 @@ def infer_quantization_format(
     if save_compressed:
         weight_args, input_args = _get_unique_quant_args(model)
         is_24_structure = (
-            SparsityStructure(sparsity_structure).value
-            == SparsityStructure.TWO_FOUR.value
+            SparsityStructure(sparsity_structure) == SparsityStructure.TWO_FOUR
         )
         is_weight_only = len(input_args) == 0 and len(weight_args) > 0
 

diff --git a/src/llmcompressor/transformers/compression/sparsity_config.py b/src/llmcompressor/transformers/compression/sparsity_config.py
@@ -1,9 +1,6 @@
 from typing import Dict, Optional
 
-from compressed_tensors import (
-    CompressionFormat,
-    SparsityCompressionConfig,
-)
+from compressed_tensors import CompressionFormat, SparsityCompressionConfig
 from compressed_tensors.config import SparsityStructure
 from torch import Tensor
 from torch.nn import Module