def InvSqrt(optimizer, warmup_step):
    """Originally used for Transformer (in "Attention is All You Need")."""

    def lr_lambda(step):
        # return a multiplier instead of a learning rate
        if step == warmup_step:  # also covers the case where both are 0
            return 1.
        else:
            return 1. / (step ** 0.5) if step > warmup_step else (step + 1) / (warmup_step ** 1.5)

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)
Originally used for Transformer (in "Attention is All You Need").
analyze/convnexts4nd/src/utils/optim/schedulers.py