fix key errors

Author: airMeng

torchao/quantization/quant_api.py

@@ -681,13 +681,14 @@ def apply_int4_weight_only_quant(weight):
         eps = 1e-6
         zero_point_dtype = torch.bfloat16
 
-        assert layout in LAYOUT_TO_ZERO_POINT_DOMAIN.keys(), f"Only support layout: {LAYOUT_TO_ZERO_POINT_DOMAIN.keys()}"
+        nonlocal zero_point_domain
+        assert any(isinstance(layout, support_layout) for support_layout in LAYOUT_TO_ZERO_POINT_DOMAIN.keys()), f"Only support layout: {LAYOUT_TO_ZERO_POINT_DOMAIN.keys()}"
         if zero_point_domain is None:
             # the first value is the default one
-            zero_point_domain = LAYOUT_TO_ZERO_POINT_DOMAIN[layout][0]
-            preserve_zero = LAYOUT_TO_PRESERVE_ZEROS[layout]
+            zero_point_domain = LAYOUT_TO_ZERO_POINT_DOMAIN[type(layout)][0]
+            preserve_zero = LAYOUT_TO_PRESERVE_ZEROS[type(layout)]
         else:
-            assert zero_point_domain in LAYOUT_TO_ZERO_POINT_DOMAIN[layout], f"Layout only support {LAYOUT_TO_ZERO_POINT_DOMAIN[layout]}"
+            assert zero_point_domain in LAYOUT_TO_ZERO_POINT_DOMAIN[type(layout)], f"Layout only support {LAYOUT_TO_ZERO_POINT_DOMAIN[layout]}"
 
         # Sparse Marlin only supports symmetric quantization.
         # NOTE: If we start having lots of layouts that require different configurations,