import tensorflow as tf

        dis_train_op = while_loop(cond, body, inputs)

        # tf.contrib.gan's train op does not manage global steps in it
        train_op = tf.group(
            dis_train_op,
            gan_train_ops.generator_train_op,
            global_step.assign_add(1))

    if params['use_tpu']:
        # TPU version of EstimatorSpec
        return tf.contrib.tpu.TPUEstimatorSpec(
            mode=mode,
            predictions=predictions,
            loss=loss,
            train_op=train_op,)
    else:
        return tf.estimator.EstimatorSpec(
            mode=mode,
            predictions=predictions,
            loss=loss,
