Missed intialize entry point with conditional unwrap

Author: Jack-Khuu

build/builder.py

@@ -440,7 +440,7 @@ def _initialize_model(
     quantize,
     tokenizer=None,
     max_seq_length=None,
-    support_tensor_subclass:bool=True,
+    support_tensor_subclass: bool = True,
 ):
     print("Loading model...")
 
@@ -511,7 +511,13 @@ def _initialize_model(
         if quantize:
             print(f"Quantizing the model with: {quantize}")
             with measure_time("Time to quantize model: {time:.02f} seconds"):
-                quantize_model(model, builder_args.device, quantize, tokenizer, support_tensor_subclass)
+                quantize_model(
+                    model,
+                    builder_args.device,
+                    quantize,
+                    tokenizer,
+                    support_tensor_subclass,
+                )
                 device_sync(device=builder_args.device)
 
         if builder_args.setup_caches:

export.py

@@ -126,6 +126,7 @@ def main(args):
             quantize,
             tokenizer,
             max_seq_length=builder_args.max_seq_length,
+            support_tensor_subclass=output_dso_path is None,
         )
         model_to_pte = model
         model_to_dso = model

quantization/quantize.py

@@ -43,18 +43,19 @@
     Int8DynActInt4WeightQuantizer,
     quantize_,
 )
+from torchao.utils import unwrap_tensor_subclass
 
 
 #########################################################################
 ###                  torchchat quantization API                       ###
 
 
 def quantize_model(
-    model: nn.Module, 
-    device, 
-    quantize_options, 
-    tokenizer=None, 
-    support_tensor_subclass:bool=True
+    model: nn.Module,
+    device,
+    quantize_options,
+    tokenizer=None,
+    support_tensor_subclass: bool = True,
 ):
     """
     Quantize the specified model using the quantizers described by