Return conversion count

neuralmagic · May 30, 2023 · 4f7246e · 4f7246e
1 parent acf3ce8
commit 4f7246e
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/sparseml/pytorch/sparsification/quantization/quantize_qat_export.py b/src/sparseml/pytorch/sparsification/quantization/quantize_qat_export.py
@@ -841,9 +841,12 @@ def _convert_quantizable_matmul_with_quantized_outputs(model: ModelProto):
 
     if matmul_nodes:
         _LOGGER.info(
-            f"Converted {conversion_count} quantizable MatMul ops " "to QLinearMatMul"
+            f"Converted {conversion_count} quantizable MatMul with quantized outputs "
+            "to QLinearMatMul"
         )
 
+    return conversion_count
+
 
 def _convert_quantizable_matmul(model: ModelProto):
     conversion_count = _convert_quantizable_matmul_with_quantized_outputs(model)