def compute_l2_reg(model,model_name):
    # Function that sets weight_decay only for weights and not biases and only 
    # for conv layers inside residual layers
    lambda_ = FLAGS.weight_decay
    params_dict = dict(model.named_parameters())
    l2_reg=[]  
    if model_name == 'resnet8':
        for key, value in params_dict.items():
            if ((key[-8:] == '2.weight' or key[-8:] == '5.weight') and key[0:8]=='residual'):
                l2_reg += [lambda_*torch.norm(value.view(value.size(0),-1),2)]
    else:
        for key, value in params_dict.items():
            if ((key[-8:] == '2.weight' or key[-8:] == '6.weight') and key[0:8]=='residual'):
                l2_reg += [lambda_*torch.norm(value.view(value.size(0),-1),2)]
    l2_reg = sum(l2_reg)
    return l2_reg

This is how I compute l2 regularization to emulate same kernel regularization as in Keras. Then I sum l2_reg to mse loss (it’s a regression problem):
loss = F.mse_loss(outputs, targets) + l2_reg