import torch
from typing import List


def set_grads(model: torch.nn.Module, trainable_layers: List[str]):
    """
    Method that enables or disables gradients of model parameters according to specified layers.

    Parameters
    ----------
    model: torch.nn.Module
        Torch model for which parameter gradients should be set
    trainable_layers: List
        List of strings, i.e. layer / parameter names, for which training is enabled. For model parameters, which do not
        match any pattern specified in trainable_layers, training is disable by setting requires_grad to False.
    """
    def is_trainable(x, trainable_layers):
        return any([(layer in x) or ('fc' in x) for layer in trainable_layers])

    for p in model.parameters():
        p.requires_grad = False

    trainable_parameters = [n for n, p in model.named_parameters() if is_trainable(n, trainable_layers)]
    for n, p in model.named_parameters():
        if n in trainable_parameters:
            p.requires_grad = True