import tensorflow as tf

        current_image_id_ph,
        stable_stage_num_images,
        transition_stage_num_images,
        num_blocks=3)
    x = tf.random_normal([2, 16, 16, 3])
    logits, _ = networks.discriminator(
        x, progress, _num_filters_stub,
        networks.ResolutionSchedule(
            start_resolutions=(4, 4), scale_base=2, num_resolutions=3))
    fake_loss = tf.reduce_sum(tf.square(logits))
    grad_norms = [
        _get_grad_norm(
            fake_loss, tf.trainable_variables('.*/progressive_gan_block_1/.*')),
        _get_grad_norm(
            fake_loss, tf.trainable_variables('.*/progressive_gan_block_2/.*')),
        _get_grad_norm(
            fake_loss, tf.trainable_variables('.*/progressive_gan_block_3/.*'))
    ]

    grad_norms_output = None
    with self.test_session(use_gpu=True) as sess:
      sess.run(tf.global_variables_initializer())
      grad_norms_output = np.array([
          sess.run(grad_norms, feed_dict={current_image_id_ph: i})
