Update learned_round.py

Xilinx · Dec 2, 2024 · e5bc47c · e5bc47c
1 parent 1a2863b
commit e5bc47c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/brevitas/core/function_wrapper/learned_round.py b/src/brevitas/core/function_wrapper/learned_round.py
@@ -76,12 +76,14 @@ def __init__(self) -> None:
         super(LearnedRoundIdentity, self).__init__()
         self.tensor_clamp = TensorClampSte()
 
+    @brevitas.jit.ignore
     def forward(self, p: torch.Tensor) -> torch.Tensor:
         return self.tensor_clamp(
             p,
             min_val=torch.tensor(-0.5, device=p.device),
             max_val=torch.tensor(+0.5, device=p.device))
 
+    @brevitas.jit.ignore
     def round_forward(self, x: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
         return round_ste(x + p)