def prune_model(
    model: Module,
    pruning_fn: Union[Callable, str],
    amount: Union[float, int],
    keys_to_prune: Optional[List[str]] = None,
    layers_to_prune: Optional[List[str]] = None,
    dim: int = None,
    l_norm: int = None,
) -> None:
    """
    Prune model function can be used for pruning certain
    tensors in model layers.

    Args:
        model: Model to be pruned.
        pruning_fn: Pruning function with API same as in torch.nn.utils.pruning.
            pruning_fn(module, name, amount).
        keys_to_prune: list of strings.
            Determines which tensor in modules will be pruned.
        amount: quantity of parameters to prune.
            If float, should be between 0.0 and 1.0 and
            represent the fraction of parameters to prune.
            If int, it represents the absolute number
            of parameters to prune.
        layers_to_prune: list of strings - module names to be pruned.
            If None provided then will try to prune every module in model.
        dim (int, optional): if you are using structured pruning method you need
            to specify dimension. Defaults to None.
        l_norm (int, optional): if you are using
            ln_structured you need to specify l_norm. Defaults to None.

    Example:
        .. code-block:: python

           pruned_model = prune_model(model, pruning_fn="l1_unstructured")

    Raises:
        AttributeError: If layers_to_prune is not None, but there is
            no layers with specified name. OR
        ValueError: if no layers have specified keys.
    """
    nn_model = get_nn_from_ddp_module(model)
    pruning_fn = get_pruning_fn(pruning_fn, l_norm=l_norm, dim=dim)
    keys_to_prune = keys_to_prune or ["weight"]
    pruned_modules = 0
    for name, module in nn_model.named_modules():
        try:
            if layers_to_prune is None or name in layers_to_prune:
                for key in keys_to_prune:
                    pruning_fn(module, name=key, amount=amount)
                pruned_modules += 1
        except AttributeError as e:
            if layers_to_prune is not None:
                raise e

    if pruned_modules == 0:
        raise ValueError(f"There is no {keys_to_prune} key in your model")

    Prune model function can be used for pruning certain
    tensors in model layers.

    Args:
        model: Model to be pruned.
        pruning_fn: Pruning function with API same as in torch.nn.utils.pruning.
            pruning_fn(module, name, amount).
        keys_to_prune: list of strings.
            Determines which tensor in modules will be pruned.
        amount: quantity of parameters to prune.
            If float, should be between 0.0 and 1.0 and
            represent the fraction of parameters to prune.
            If int, it represents the absolute number
            of parameters to prune.
        layers_to_prune: list of strings - module names to be pruned.
            If None provided then will try to prune every module in model.
        dim (int, optional): if you are using structured pruning method you need
            to specify dimension. Defaults to None.
        l_norm (int, optional): if you are using
            ln_structured you need to specify l_norm. Defaults to None.

    Example:
        .. code-block:: python

           pruned_model = prune_model(model, pruning_fn="l1_unstructured")

    Raises:
        AttributeError: If layers_to_prune is not None, but there is
            no layers with specified name. OR
        ValueError: if no layers have specified keys.
    
catalyst/utils/pruning.py