Please refer to the code of optimizer in PyTorch. In detail, after backward, the weight will be added to the grad of weight~(L2 weight decay). We could also directly use the above solution to avoid apply weight decay to bn. However, I have another more elegant method like function below:

def apply_weight_decay(*modules, weight_decay_factor=0., wo_bn=True):
    '''
    https://discuss.pytorch.org/t/weight-decay-in-the-optimizers-is-a-bad-idea-especially-with-batchnorm/16994/5
    Apply weight decay to pytorch model without BN;
    In pytorch:
        if group['weight_decay'] != 0:
            grad = grad.add(p, alpha=group['weight_decay'])
    p is the param;
    :param modules:
    :param weight_decay_factor:
    :return:
    '''
    for module in modules:
        for m in module.modules():
            if hasattr(m, 'weight'):
                if wo_bn and isinstance(m, torch.nn.modules.batchnorm._BatchNorm):
                    continue
                m.weight.grad += m.weight * weight_decay_factor

Please note that it should be applied after loss.backward. Besides, as far as I know, all bn layers in pytorch are inherited from  _BatchNorm class.