def get_lr_schedule(lr_schedule):
    boundaries = lr_schedule.boundaries
    values = lr_schedule.values
    return tf.keras.optimizers.schedules.PiecewiseConstantDecay(boundaries, values)
def get_lr_schedule(lr_schedule):
    boundaries = lr_schedule.boundaries
    values = lr_schedule.values
    return tf.keras.optimizers.schedules.PiecewiseConstantDecay(boundaries, values)