How can I add custom regularization to my loss?
I use cross-entropy with weights:
criterion = torch.nn.CrossEntropyLoss(weight=weights)
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.