import tensorflow as tf

      p = self._testParams()
      tp = p.train
      tp.lr_schedule.boundaries = [300000, 400000, 500000]
      tp.lr_schedule.values = [1.0, 0.1, 0.01, 0.001]
      lrs = tp.lr_schedule.Instantiate()
      steps = [299999, 300001, 399999, 400001, 499999, 500001]
      fetches = [lrs.Value(_) for _ in steps]
      values = sess.run(fetches)
      self.assertAllClose([1.0, 0.1, 0.1, 0.01, 0.01, 0.001], values)

  def testBatchSplit(self):

    def Run(num_splits):
      p = self._testParams()
      with self.session(use_gpu=False, graph=tf.Graph()) as sess:
        tf.set_random_seed(93820981)
        p.is_eval = True
        p.input.cur_iter_in_seed = False
        p.input.bucket_batch_limit = [
            b * 2 / num_splits for b in p.input.bucket_batch_limit
        ]
        with cluster_factory.ForTestingWorker(gpus=num_splits):
          mdl = p.Instantiate()
          metrics = mdl.FPropDefaultTheta()[0]
        tf.global_variables_initializer().run()
        return sess.run(metrics['loss'])

    res1, res2 = Run(1), Run(2)
    self.assertAllClose(res1[0], res2[0])
    self.assertAllEqual(res1[1], res2[1])
