Keras callback function TypeError: unsupported operand type(s)

Question

I'm creating a text classification model using RoBERTa model. I keep encountering this error TypeError: unsupported operand type(s) for *: 'WarmUp' and 'int' whenever I use either ReduceLROnPlateau or LearningRateScheduler in my callback function.

This is my code:

epochs = 30
steps_per_epoch = tf.data.experimental.cardinality(train_ds).numpy()
num_train_steps = steps_per_epoch * epochs
num_warmup_steps = int(0.1 * num_train_steps)
init_lr = 3e-5

callback = [tf.keras.callbacks.EarlyStopping(monitor='val_accuracy', 
                                             min_delta=0,
                                             patience=3,
                                             verbose=1,
                                             mode='auto',
                                             baseline=None,
                                             restore_best_weights=False,
                                             start_from_epoch=0),
           tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', 
                                                factor=0.2,
                                                patience=5, 
                                                min_lr=0.001)
           ]

optimizer = optimization.create_optimizer(init_lr=init_lr,
                                          num_train_steps=num_train_steps,
                                          num_warmup_steps=num_warmup_steps,
                                          optimizer_type='adamw')

classifier_model.compile(optimizer=optimizer,
                         loss=loss,
                         metrics=metrics)

print(f'Training model with {tfhub_handle_encoder}')

history = classifier_model.fit(x=train_ds,
                                   validation_data=val_ds,
                                   epochs=epochs,
                                   callbacks=callback,
                                   steps_per_epoch=steps_per_epoch,
                                   verbose=1)

This is the full error message:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[63], line 35
     29 classifier_model.compile(optimizer=optimizer,
     30                          loss=loss,
     31                          metrics=metrics)
     33 print(f'Training model with {tfhub_handle_encoder}')
---> 35 history = classifier_model.fit(x=train_ds,
     36                                    validation_data=val_ds,
     37                                    epochs=epochs,
     38                                    callbacks=callback,
     39                                    steps_per_epoch=steps_per_epoch,
     40                                    verbose=1)

File /usr/local/lib/python3.8/site-packages/keras/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File /usr/local/lib/python3.8/site-packages/keras/utils/generic_utils.py:210, in Progbar.update(self, current, values, finalize)
    208 value_base = max(current - self._seen_so_far, 1)
    209 if k not in self._values:
--> 210     self._values[k] = [v * value_base, value_base]
    211 else:
    212     self._values[k][0] += v * value_base

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

I'm quite new to this so I'm clueless. I've no idea where the WarmUp is from. I think it shouldn't be the num_warmup_steps since I already typecasted it as int.