Update Error check

Author: Jack-Khuu

torchchat/utils/quantize.py

@@ -52,7 +52,7 @@
 
 
 # Flag for whether the a8wxdq quantizer is available.
-a8wxdq_loaded = False
+a8wxdq_load_error: Optional[Exception] = None
 
 #########################################################################
 ###                  torchchat quantization API                       ###
@@ -79,6 +79,10 @@ def quantize_model(
         quantize_options = json.loads(quantize_options)
 
     for quantizer, q_kwargs in quantize_options.items():
+        # Test if a8wxdq quantizer is available; Surface error if not.
+        if quantizer == "linear:a8wxdq" and a8wxdq_load_error is not None:
+            raise Exception(f"Note: Failed to load torchao experimental a8wxdq quantizer with error: {a8wxdq_load_error}")
+
         if (
             quantizer not in quantizer_class_dict
             and quantizer not in ao_quantizer_class_dict
@@ -100,9 +104,6 @@ def quantize_model(
 
             try:
                 if quantizer == "linear:a8wxdq":
-                    if not a8wxdq_loaded:
-                        raise Exception(f"Note: Failed to load torchao experimental a8wxdq quantizer with error: {e}")
-
                     quant_handler = ao_quantizer_class_dict[quantizer](
                         device=device,
                         precision=precision,
@@ -904,8 +905,5 @@ def quantized_model(self) -> nn.Module:
         print("Failed to torchao ops library with error: ", e)
         print("Slow fallback kernels will be used.")
 
-    # Mark the Quant option as available
-    a8wxdq_loaded = True
-
 except Exception as e:
-    pass
+    a8wxdq_load_error = e