Feat (quantize): add support for Identity

Xilinx · Jul 13, 2023 · f1c9f27 · f1c9f27
1 parent f93911e
commit f1c9f27
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/brevitas/graph/quantize_impl.py b/src/brevitas/graph/quantize_impl.py
@@ -32,7 +32,8 @@
     nn.AdaptiveAvgPool2d,
     nn.AdaptiveAvgPool3d,
     nn.PixelShuffle,
-    nn.PixelUnshuffle)
+    nn.PixelUnshuffle,
+    nn.Identity)
 
 PRECISION_PRESERVING_MODULES = (
     nn.Dropout,
@@ -42,7 +43,8 @@
     nn.MaxPool2d,
     nn.MaxPool3d,
     nn.PixelShuffle,
-    nn.PixelUnshuffle)
+    nn.PixelUnshuffle,
+    nn.Identity)
 
 
 def inp_placeholder_handler(model, input_quantizer):