Add costum regularization to loss

Your current code snippet would use a weighted loss (which might be seen as a type of regularization).
The usual work flow would be to add the regularization term directly to the loss as seen here.