def testRuns(self):
    with mock.patch.object(dataset, 'load', new=fake_mnist_data):
      train.train(
          batch_size=1,
          learning_rate=0.1,
          num_training_iters=10,
          validation_steps=5)
      train.train(
          batch_size=2,
          learning_rate=0.1,
          num_training_iters=5,
          validation_steps=2)
      train.train(
          batch_size=10,
          learning_rate=0.1,
          num_training_iters=1,
          validation_steps=1)