Skip to content

Commit c6136b2

Browse files
committed
update
1 parent 246d711 commit c6136b2

File tree

2 files changed

+6
-7
lines changed

2 files changed

+6
-7
lines changed

src/compressed_tensors/compressors/model_compressors/model_compressor.py

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,7 @@ def from_pretrained_model(
185185
# assume multiple compression formats means mixed-precision
186186
# as we currently only support one compressor per precision type and scheme
187187
if len(quantization_format) > 1:
188-
quantization_format = CompressionFormat.mixed_precision
188+
quantization_format = CompressionFormat.mixed_precision.value
189189
else:
190190
quantization_format = quantization_format[0]
191191

@@ -269,15 +269,15 @@ def parse_quantization_config(
269269

270270
return quantization_config
271271

272-
def _fetch_unique_quantization_formats(self):
272+
def _fetch_unique_quantization_formats(self) -> List[str]:
273273
"""
274-
Get all unique compression formats used in
275-
model
274+
Get all unique compression formats present in a model
275+
:return: list of quantization formats
276276
"""
277277
quantization_formats = []
278278
for _, scheme in self.quantization_config.config_groups.items():
279279
if scheme.format not in quantization_formats:
280-
quantization_formats.append(scheme)
280+
quantization_formats.append(scheme.format)
281281
return quantization_formats
282282

283283
def __init__(
@@ -302,7 +302,6 @@ def __init__(
302302
)
303303

304304
quantization_formats = self._fetch_unique_quantization_formats()
305-
306305
if quantization_config is not None:
307306
self.quantization_compressor = {}
308307
for format in quantization_formats:

src/compressed_tensors/quantization/quant_scheme.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ class QuantizationScheme(BaseModel):
4949
weights: Optional[QuantizationArgs] = None
5050
input_activations: Optional[QuantizationArgs] = None
5151
output_activations: Optional[QuantizationArgs] = None
52-
format: Optional[CompressionFormat] = None
52+
format: Optional[str] = None
5353

5454
@model_validator(mode="after")
5555
def validate_model_after(model: "QuantizationScheme") -> "QuantizationScheme":

0 commit comments

Comments
 (0)