From 1abc765c43bd30b1e73e40cbde0a7edd5e5c5ec1 Mon Sep 17 00:00:00 2001 From: Giuseppe Franco Date: Thu, 13 Jul 2023 20:39:12 +0100 Subject: [PATCH] Feat (equalize): add support for Identity --- src/brevitas/graph/equalize.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/brevitas/graph/equalize.py b/src/brevitas/graph/equalize.py index bb28c0234..207878001 100644 --- 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__)