import torch
from typing import List
from typing import Dict
from typing import Any


def add_weight_decay(
    model: torch.nn.Module, weight_decay: float
) -> List[Dict[str, Any]]:
    """Add weight decay to only the convolutional kernels. Do not add weight decay
    to biases and BN. TensorFlow does this automatically, but for some reason this is
    the PyTorch default.

    Args:
        model: The model where the weight decay is applied.
        weight_decay: How much decay to apply.

    Returns:
        A list of dictionaries encapsulating which params need weight decay.
    """
    decay, no_decay = [], []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if len(param.shape) == 1:
            no_decay.append(param)
        else:
            decay.append(param)

    return [
        {"params": no_decay, "weight_decay": 0.0},
        {"params": decay, "weight_decay": weight_decay},
    ]