diff --git a/src/compressed_tensors/compressors/model_compressors/model_compressor.py b/src/compressed_tensors/compressors/model_compressors/model_compressor.py index 24dcfe95..87050cfd 100644 --- a/src/compressed_tensors/compressors/model_compressors/model_compressor.py +++ b/src/compressed_tensors/compressors/model_compressors/model_compressor.py @@ -207,7 +207,7 @@ def from_pretrained_model( @staticmethod def parse_sparsity_config( - compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"] + compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"], ) -> Union[Dict[str, Any], None]: """ Parse sparsity config from quantization/compression config. Sparsity @@ -227,7 +227,7 @@ def parse_sparsity_config( @staticmethod def parse_quantization_config( - compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"] + compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"], ) -> Union[Dict[str, Any], None]: """ Parse quantization config from quantization/compression config. The diff --git a/tests/test_transform/factory/test_memory.py b/tests/test_transform/factory/test_memory.py index 64a068c9..34db4ed5 100644 --- a/tests/test_transform/factory/test_memory.py +++ b/tests/test_transform/factory/test_memory.py @@ -85,8 +85,5 @@ def test_memory_sharing(type, randomize, requires_grad, offload=False): @requires_accelerate() @pytest.mark.parametrize("type", ("hadamard", "random-hadamard")) @pytest.mark.parametrize("randomize", (True, False)) -def test_memory_sharing_offload( - type, - randomize, -): +def test_memory_sharing_offload(type, randomize): test_memory_sharing(type, randomize, requires_grad=False, offload=True)