Can you please refer to this and let me know if you are inheriting from the correct class of tf.keras.callback.Callback ?