diff --git a/distiller/pruning/ranked_structures_pruner.py b/distiller/pruning/ranked_structures_pruner.py index d6d3fe16c8d09564138d2622facb0c77969c840c..0ae0c1211183178b6a874f6d28905c7b6ea216f8 100755 --- a/distiller/pruning/ranked_structures_pruner.py +++ b/distiller/pruning/ranked_structures_pruner.py @@ -119,7 +119,8 @@ class LpRankedStructureParameterPruner(RankedStructureParameterPruner): group_pruning_fn = partial(self.rank_and_prune_blocks, block_shape=self.block_shape) binary_map = group_pruning_fn(fraction_to_prune, param, param_name, - zeros_mask_dict, model, binary_map, self.magnitude_fn) + zeros_mask_dict, model, binary_map, + magnitude_fn=self.magnitude_fn) return binary_map @staticmethod