import torch


def l_n_reg(model, device, norm = 2):
    """
    Wrapper to add a weight regularizer, with adjustable norm.
    ----------
    model: model with the weights to be regularized
    device: either cpu or gpu
    norm: default 2, the norm to use in regularization
    """
    reg_loss = torch.tensor(0., requires_grad=True).to(device)
    for name, param in model.named_parameters():
        if 'weight' in name:
            reg_loss = reg_loss + torch.norm(param, norm)
    return reg_loss