Thanks, if I want to output the penalty, for example l2 loss. Do you think the following workable, as a simple implemetation?
def get_reg_loss(model):
     reg_loss = 0

     for param in model.parameters():
         reg_loss += param**2

    _lambda = 0.001
    reg_loss += _lambda * reg_loss
    return reg_loss