import tensorflow as tf


        gan_train_ops = tf.contrib.gan.gan_train_ops(gan_model, gan_loss, gen_optimizer, dis_optimizer)

        while_loop = tf.contrib.tpu.while_loop if params['use_tpu'] else tf.while_loop

        # train the discriminator 100 steps
        inputs = [tf.constant(0), tf.constant(0.0)]
        cond = lambda i, x: tf.less(i, 100)
        def body(i, x):
            return tf.add(i, 1), gan_train_ops.discriminator_train_op

        dis_train_op = while_loop(cond, body, inputs)
