Update calibrate.py

Xilinx · Jun 23, 2023 · 31a401b · 31a401b
1 parent 6f9a5ba
commit 31a401b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/brevitas/graph/calibrate.py b/src/brevitas/graph/calibrate.py
@@ -138,7 +138,7 @@ def disable_act_quant_hook(self, module, inp, output):
     def disable_act_quantization(self, model, is_training, collect_stats=False):
         for module in model.modules():
             if isinstance(module, ActQuantProxyFromInjector):
-                module.disable_quant = collect_stats
+                module.disable_quant = not collect_stats
                 hook = module.register_forward_hook(self.disable_act_quant_hook)
                 module.train(is_training)
                 self.disable_act_quant_hooks.append(hook)