fix the bug in line no.336 in mask_conflict.py (microsoft#3629)

Co-authored-by: Xuesong Wang <wangxuesong@dm-ai.cn>

nni/compression/pytorch/utils/mask_conflict.py

@@ -333,7 +333,7 @@ def fix_mask(self):
                 elif type(m).__name__ == 'Linear':
                     new_mask[:, merged_index] = 1.
                 elif type(m).__name__ == 'BatchNorm2d':
-                    new_mask = merged_index.type_as(orig_mask)
+                    new_mask = merged_channel_mask.type_as(orig_mask)
                 else:
                     raise RuntimeError(
                         f'unsupported module type: {type(m).__name__}')