From 4db00bc0dfaee0df2be11bedc74f6f261d3b2f3e Mon Sep 17 00:00:00 2001 From: i-colbert Date: Tue, 23 Jan 2024 08:57:43 -0800 Subject: [PATCH] Fix (a2q+): adding JIT decorator --- src/brevitas/core/zero_point.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/brevitas/core/zero_point.py b/src/brevitas/core/zero_point.py index 633918d1d..193707dc0 100644 --- 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)