def METHOD_NAME(self, update_step):
    """
    return stage, and the corresponding steps within the current stage
    """
    if update_step < self.warmup_steps:
        # warmup state
        return 0, update_step
    offset = self.warmup_steps
    if update_step < offset + self.hold_steps:
        # hold stage
        return 1, update_step - offset
    offset += self.hold_steps
    if update_step <= offset + self.decay_steps:
        # decay stage
        return 2, update_step - offset
    offset += self.decay_steps
    # still here ? constant lr stage
    return 3, update_step - offset