import torch


def should_decay_parameter(name: str, param: torch.Tensor) -> bool:
    """Determines if L2 (or L2-SP) decay should be applied to parameter

    BatchNorm and bias parameters are excluded. Uses both the name of the parameter and shape of parameter to determine

    Helpful source:
    https://github.com/rwightman/pytorch-image-models/blob/198f6ea0f3dae13f041f3ea5880dd79089b60d61/timm/optim/optim_factory.py
    
    Parameters
    ----------
    name : str
        name of parameter
    param : torch.Tensor
        parameter

    Returns
    -------
    bool
        Whether or not to decay
    """
    
    if not param.requires_grad:
        return False
    elif 'batchnorm' in name.lower() or 'bn' in name.lower() or 'bias' in name.lower():
        return False
    elif param.ndim == 1:
        return False
    else:
        return True