post-merge fixes

Signed-off-by: Brian Dellabetta <[email protected]>

Author: brian-dellabetta

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -754,8 +754,8 @@ def map_module_to_scheme(model: Module) -> Dict[str, QuantizationScheme]:
         fix_fsdp_module_name(name): module.quantization_scheme
         for name, module in model.named_modules()
         if (
-            hasattr(module, "quantization_scheme") and
-            module.quantization_scheme.weights is not None
+            hasattr(module, "quantization_scheme")
+            and module.quantization_scheme.weights is not None
         )
     }
 

src/compressed_tensors/transform/factory/hadamard.py

@@ -15,7 +15,6 @@
 import math
 from typing import Optional
 
-import math
 import torch
 from compressed_tensors.transform import TransformArgs, TransformScheme
 from compressed_tensors.transform.factory.base import TransformBase, TransformFactory
@@ -24,10 +23,7 @@
     apply_transform_weight,
     get_transform_size,
 )
-from compressed_tensors.utils import (
-    get_execution_device,
-    get_offloaded_device,
-)
+from compressed_tensors.utils import get_execution_device, get_offloaded_device
 from compressed_tensors.utils.helpers import ParameterizedDefaultDict
 from torch import Tensor, device, dtype
 from torch.nn import Module, Parameter
@@ -107,7 +103,8 @@ def forward(self, value: Tensor) -> Tensor:
 
         if self.args.inverse:
             weight = weight.T
- 
-        return apply_transform_weight(
-            weight, value, self.args.location, self.module_type
-        ) / self._scale
+
+        return (
+            apply_transform_weight(weight, value, self.args.location, self.module_type)
+            / self._scale
+        )

tests/test_transform/factory/test_correctness.py

@@ -91,14 +91,17 @@ def test_correctness_embedding(type, randomized, embed_loc, linear_loc):
 
 @pytest.mark.parametrize("type", ("hadamard", "random-hadamard"))
 @pytest.mark.parametrize("randomize", (True, False))
-def test_correctness_model(type, randomize, model_apply, offload=False):
+@pytest.mark.parametrize("input_batch_size", (1, 5, 17))
+def test_correctness_model(
+    type, randomize, input_batch_size, model_apply, offload=False
+):
     # load model
     model = model_apply[0]
     if offload:
         model = offloaded_dispatch(model, torch.device("cuda"))
 
     # get output
-    input = torch.rand((17, 5, model.fcs[0].in_features))
+    input = torch.rand((input_bathc_size, 5, model.fcs[0].in_features))
     if offload:
         input = input.to(torch.device("cuda"))
     true_output = model(input)