A pytorch port of google-research/google-research/robust_loss/ - View it on GitHub