Fix (a2q+): adding JIT decorator

Xilinx · Jan 23, 2024 · 4db00bc · 4db00bc
1 parent dafedb9
commit 4db00bc
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/brevitas/core/zero_point.py b/src/brevitas/core/zero_point.py
@@ -311,6 +311,7 @@ def __init__(
         self.stats_output_shape = pre_zero_point_shape
         self.stats_input_view_shape_impl = pre_zero_point_stats_input_view_shape_impl
 
+    @brevitas.jit.script_method
     def get_zero_center(self, x: Tensor) -> Tensor:
         x = self.stats_input_view_shape_impl(x)
         u = torch.mean(x, axis=self.stats_reduce_dim, keepdim=True)