Fixes

Xilinx · Sep 11, 2024 · 9d82357 · 9d82357
1 parent 3ca7505
commit 9d82357
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 8 deletions.
diff --git a/src/brevitas/export/inference/handler.py b/src/brevitas/export/inference/handler.py
@@ -68,10 +68,9 @@ class IntWeightInferencetHandler(IntInferencetHandler):
     def prepare_for_export(self, module):
         if module.is_quant_enabled:
             self.cached_weight = None
+            super().prepare_for_export(module)
             if module._cached_weight is not None and not module.cache_inference_quant_weight_metadata_only:
-                self.cached_weight = module._cached_weight
-            else:
-                super().prepare_for_export(module)
+                self.cached_weight = module._cached_weight.value
 
     def forward(self, x) -> Tuple[torch.Tensor]:
         if self.cached_weight is not None:
@@ -142,10 +141,9 @@ class FloatWeightInferencetHandler(FloatInferencetHandler):
     def prepare_for_export(self, module):
         if module.is_quant_enabled:
             self.cached_weight = None
+            super().prepare_for_export(module)
             if module._cached_weight is not None and not module.cache_inference_quant_weight_metadata_only:
-                self.cached_weight = module._cached_weight
-            else:
-                super().prepare_for_export(module)
+                self.cached_weight = module._cached_weight.value
 
     def forward(self, x) -> Tuple[torch.Tensor]:
         if self.cached_weight is not None:

diff --git a/src/brevitas_examples/imagenet_classification/ptq/ptq_evaluate.py b/src/brevitas_examples/imagenet_classification/ptq/ptq_evaluate.py
@@ -488,8 +488,7 @@ def main():
         device, dtype = param.device, param.dtype
         ref_input = generate_ref_input(args, device, dtype)
         quant_model(ref_input)
-        compiled_model = torch.compile(
-            quant_model, fullgraph=True, dynamic=True, disable=not args.compile)
+        compiled_model = torch.compile(quant_model, fullgraph=True, disable=not args.compile)
         validate(val_loader, compiled_model, stable=dtype != torch.bfloat16)
 
     if args.export_onnx_qcdq or args.export_torch_qcdq: