def noam_decay(step, warmup_steps, model_size):
    """Learning rate schedule described in
    https://arxiv.org/pdf/1706.03762.pdf.
    """
    return (
        model_size ** (-0.5) *
        min(step ** (-0.5), step * warmup_steps**(-1.5)))
Learning rate schedule described in
    https://arxiv.org/pdf/1706.03762.pdf.
    
lsp_model/optim.py