4e21e721a65d0ac7c8465c6b7842dd39338bb3d0,src/sdk/pynni/nni/compression/torch/pruners.py,LevelPruner,calc_mask,#LevelPruner#Any#Any#,32

Before Change

            self.mask_calculated_ops.add(op_name)
        else:
            assert op_name in self.mask_dict, "op_name not in the mask_dict"
            mask = self.mask_dict[op_name]
        return mask


class AGP_Pruner(Pruner):

After Change

            threshold = torch.topk(w_abs.view(-1), k, largest=False)[0].max()
            mask_weight = torch.gt(w_abs, threshold).type_as(weight)
            mask = {"weight": mask_weight}
            if_calculated.copy_(torch.tensor(1)) // pylint: disable=not-callable
            return mask
        else:
            return None

Instances