[quantization] fix cle convT fused_bn not clear (#382)

alibaba · Dec 19, 2024 · fac2af5 · fac2af5
1 parent f897615
commit fac2af5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tinynn/graph/quantization/algorithm/cross_layer_equalization.py b/tinynn/graph/quantization/algorithm/cross_layer_equalization.py
@@ -354,5 +354,5 @@ def model_rewrite(model, dummy_input, work_dir='out') -> nn.Module:
 def clear_model_fused_bn(model: nn.Module):
     """remove the attached bn from fused conv"""
     for mod in model.modules():
-        if isinstance(mod, nn.Conv2d) and hasattr(mod, 'fused_bn_'):
+        if isinstance(mod, (nn.Conv2d, nn.ConvTranspose2d)) and hasattr(mod, 'fused_bn_'):
             delattr(mod, 'fused_bn_')