LAMB optimizer fails to work with tf.keras.optimizers.schedules.PolynomialDecay

System information

OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 11
TensorFlow version and how it was installed (source or binary): 2.10 (Anaconda)
TensorFlow-Addons version and how it was installed (source or binary): 0.19.0 (pip)
Python version: 3.10
Is GPU used? (yes/no): yes

Describe the bug

LAMB fails to work with PolynomialDecay learning rate

File c:\Users\ma7mo\mambaforge\envs\bsc\lib\site-packages\keras\utils\generic_utils.py:965, in Progbar.update(self, current, values, finalize)
    963 value_base = max(current - self._seen_so_far, 1)
    964 if k not in self._values:
--> 965     self._values[k] = [v * value_base, value_base]
    966 else:
    967     self._values[k][0] += v * value_base

TypeError: unsupported operand type(s) for *: 'PolynomialDecay' and 'int'

Code to reproduce the issue

LR = tf.keras.optimizers.schedules.PolynomialDecay(
    1e-2,
    1000,
    end_learning_rate=0.0001,
)
model = ...
model.compile(loss=tf.keras.losses.BinaryCrossentropy(),
              metrics=['accuracy'],
              optimizer=tfa.optimizers.LAMB(LR, weight_decay=0)
              )
model.fit(...)

tensorflow / addons

LAMB optimizer fails to work with tf.keras.optimizers.schedules.PolynomialDecay #2798