Fix: disable/enable act quant

Xilinx · Jul 7, 2023 · 62b6a58 · 62b6a58
1 parent dc47ce1
commit 62b6a58
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/brevitas_examples/imagenet_classification/ptq/learned_round_utils.py b/src/brevitas_examples/imagenet_classification/ptq/learned_round_utils.py
@@ -74,6 +74,7 @@ def __call__(self, model_input):
 
         handle = self.layer.register_forward_hook(self.data_saver)
         self.disable_param_quantization(self.model, is_training=False)
+        self.disable_act_quantization(self.model, is_training=False)
         with torch.no_grad():
             try:
                 _ = self.model(model_input)
@@ -83,6 +84,7 @@ def __call__(self, model_input):
 
         self.data_saver.store_output = False
         self.enable_param_quantization(self.model, is_training=False)
+        self.enable_act_quantization(self.model, is_training=False)
         with torch.no_grad():
             try:
                 _ = self.model(model_input)