code: def on_epoch_end(self, epoch, logs=None):
      current_lr = tf.keras.backend.get_value(self.model.optimizer.lr)
      print("epoch {} current lr is {}".format(epoch, current_lr))
      if epoch < 5:
          assert abs(current_lr - self.warmup_lr[epoch]) < 1e-5
      elif 5 <= epoch < 10:
          assert abs(current_lr - 0.4) < 1e-5
      elif 10 <= epoch < 15:
          assert abs(current_lr - 0.04) < 1e-5
      elif 15 <= epoch < 20:
          assert abs(current_lr - 0.004) < 1e-5
      else:
          assert abs(current_lr - 0.0004) < 1e-5
  
imports: import os
  from unittest import TestCase
  import numpy
  import pytest
  import tensorflow
  from bigdl.orca.learn.tf2 import Estimator
  from bigdl.orca.ray import OrcaRayContext
  from bigdl.orca import OrcaContext
label_desc: Use image node_id 2 for calling the LRChecker obj's underlying member method code with example usage: obj.on_epoch_end(epoch, logs) without return types