Custom callbacks with warm up and cosine decay in TensorFlow

I referred to this example and especially the callback presented here. So, I decided to write out a callback inspired by this one. Basically, it combines warm-ups and cosine decays.

Here's how I coded it up -

class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, base_lr=0.1, end_lr=0.001, warmup_steps=390*5):
        super(CustomSchedule, self).__init__()

        self.base_lr = base_lr
        self.end_lr = end_lr
        self.warmup_steps = warmup_steps
    
    def __call__(self, step=390*35):
        warmup_lr_schedule = tf.linspace(0., self.base_lr, self.warmup_steps)
        iters = tf.range(step, dtype=tf.float32) 
        cosine_lr_schedule = tf.convert_to_tensor([self.end_lr + 0.5 * (self.base_lr - self.end_lr) * (1 + \
                        tf.math.cos(tf.constant(math.pi) * t / (step))) for t in iters])
        lr_schedule = tf.concat([warmup_lr_schedule, cosine_lr_schedule], axis=0)
        
        return lr_schedule

I verified if this is the one I wanted and indeed it is -

But when I pass this callback inside an optimizer I run into weird stuff -

OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

Here's my Colab Notebook for full reproducibility. I am aware that the LearningRateSchedule.__call__ method accepts a step argument and outputs the learning rate to use for that particular step. But the class I presented is outputting an entire schedule every step rather than the learning rate for that particular step. But I could not figure out how could achieve this effectively.