def remove(self, module: _torch.nn.Module, fuse_pruning_mask: bool = True) -> _torch.nn.Module:
        """Removes pruning masks and forward_pre_hooks from the module

        If `fuse_pruning_mask` is True, then weights are fused with the pruning
        mask before re-registering the weights under the original name
        """
        name = self._tensor_name
        for k, hook in module._forward_pre_hooks.items():
            if isinstance(hook, BaseDynamicPruningMethod) and hook._tensor_name == name:
                self._remove_impl(module, fuse_pruning_mask)
                del module._forward_pre_hooks[k]
                if hasattr(module, "pruning_method"):
                    delattr(module, "pruning_method")
                return module

        raise ValueError(
            f"Parameter '{name}' of module {module} has to be pruned "
            f"before pruning can be removed."
        )
Removes pruning masks and forward_pre_hooks from the module

        If `fuse_pruning_mask` is True, then weights are fused with the pruning
        mask before re-registering the weights under the original name
        
coremltools/optimize/torch/pruning/_base_pruning_method.py