lr_reg_wih_lbfgs.py

reg_loss = 0
for mod in self.model.modules():
    if isinstance(mod, _BatchNorm):
        if self.decay_bn:
            for name, param in mod.named_parameters(recurse=False):
                reg_loss = reg_loss + param.norm(2)    
    else:
        for name, param in mod.named_parameters(recurse=False):
            if not name.endswith("bias"):
                reg_loss = reg_loss + param.norm(2)
            else:
                if self.decay_bias:
                    reg_loss = reg_loss + param.norm(2)