import tensorflow as tf

    self.assertEqual(len(grads), len(vars_))
    for grad, var in zip(grads, vars_):
      if grad is not None:
        self.assertEqual(grad.shape, var.shape)

  def test_training_graph(self):
    """Test model training in graph mode."""
    with tf.Graph().as_default():
      config = config_.get_hparams_cifar_38()
      config.add_hparam("n_classes", 10)
      config.add_hparam("dataset", "cifar-10")

      x = tf.random_normal(
          shape=(self.config.batch_size,) + self.config.input_shape)
