Small fixes for torchao quant (#2476)

2024-12-16 14:08:12 -08:00
parent 7154b4b1df
commit 82699474fd
2 changed files with 6 additions and 5 deletions
@@ -26,11 +26,12 @@ def apply_torchao_config_to_model(
        quantize_,
    )
    from torchao.quantization.observer import PerRow, PerTensor
+    from torchao.quantization.quant_api import _is_linear

    if filter_fn is None:

        def filter_fn(module, fqn):
-            return "proj" in fqn
+            return _is_linear(module) and "proj" in fqn

    if torchao_config == "" or torchao_config is None:
        return model