[Bug] fix bug in find_make_group_parser for depthwise (#195)

Co-authored-by: yangmin <yangmin@microbt.com>
open-mmlab · Jul 8, 2022 · 8b57a07 · 8b57a07
1 parent c83a071
commit 8b57a07
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/mmrazor/models/pruners/structure_pruning.py b/mmrazor/models/pruners/structure_pruning.py
@@ -354,6 +354,8 @@ def find_make_group_parser(self, node_name, name2module):
         elif (node_name in name2module
               and isinstance(name2module[node_name], nn.Conv2d)
               and name2module[node_name].in_channels
+              == name2module[node_name].out_channels
+              and name2module[node_name].in_channels
               == name2module[node_name].groups):
             return MAKE_GROUP_PARSER_DICT['depthwise']
         else: