Feat (equalize): add support for Identity

Xilinx · Jul 13, 2023 · 1abc765 · 1abc765
1 parent f1c9f27
commit 1abc765
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/brevitas/graph/equalize.py b/src/brevitas/graph/equalize.py
@@ -56,7 +56,8 @@
  nn.AvgPool3d,
  nn.AdaptiveAvgPool1d,
  nn.AdaptiveAvgPool2d,
- nn.AdaptiveAvgPool3d)
+ nn.AdaptiveAvgPool3d,
+ nn.Identity)
 
 _scale_invariant_op = (torch.mul, operator.mul, operator.imul, operator.__mul__, operator.__imul__)