import tensorflow as tf


    generator_inputs = features
    real_data = labels

    gan_model = tf.contrib.gan.gan_model(generator_fn, discriminator_fn, real_data, generator_inputs)

    predictions = gan_model.generated_data
    loss = None
    train_op = None

    if mode == tf.estimator.ModeKeys.TRAIN:
        # define loss
        gan_loss = tf.contrib.gan.gan_loss(gan_model, add_summaries=False)
        loss = gan_loss.generator_loss

        # define train_op
        gen_optimizer = tf.train.RMSPropOptimizer(learning_rate=0.05)
        dis_optimizer = tf.train.RMSPropOptimizer(learning_rate=0.05)

        # wrapper to make the optimizer work with TPUs
        if params['use_tpu']:
            gen_optimizer = tf.contrib.tpu.CrossShardOptimizer(gen_optimizer)
            dis_optimizer = tf.contrib.tpu.CrossShardOptimizer(dis_optimizer)

