Remove non compatible JIT modules

Xilinx · Feb 1, 2024 · be0465f · be0465f
1 parent c3faec7
commit be0465f
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/brevitas_examples/common/generative/quant_blocks.py b/src/brevitas_examples/common/generative/quant_blocks.py
@@ -89,7 +89,8 @@ def forward(self, x: Tensor, scale: Tensor, bit_width: Tensor):
         return zero_point
 
 
-class RuntimeDynamicStatsScaling(brevitas.jit.ScriptModule):
+# TODO: restore JIT compatibility
+class RuntimeDynamicStatsScaling(nn.Module):
 
     def __init__(
             self,
@@ -101,7 +102,6 @@ def __init__(
         self.stats_impl = scaling_stats_impl
         self.dynamic_scaling_broadcastable_fn = dynamic_scaling_broadcastable_fn
 
-    @brevitas.jit.script_method
     def forward(self, x) -> Tensor:
         shape = x.shape
         x = self.scaling_stats_input_view_shape_impl(x)
@@ -110,7 +110,8 @@ def forward(self, x) -> Tensor:
         return x
 
 
-class RuntimeDynamicStatsZeroPoint(brevitas.jit.ScriptModule):
+# TODO: restore JIT compatibility
+class RuntimeDynamicStatsZeroPoint(nn.Module):
 
     def __init__(
             self,
@@ -125,7 +126,6 @@ def __init__(
         self.dynamic_scaling_broadcastable_fn = dynamic_scaling_broadcastable_fn
         self.scale_shift_zero_point = _ScaleShiftZeroPoint(int_quant, quantize_zero_point)
 
-    @brevitas.jit.script_method
     def forward(self, x, scale, bit_width) -> Tensor:
         shape = x.shape
         x = self.zero_point_stats_input_view_shape_impl(x)