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__)