import tensorflow as tf

    if FLAGS.allow_mix_precision:
        custom_op.parameter_map["precision_mode"].s = tf.compat.as_bytes("allow_mix_precision")
    if FLAGS.auto_tune:
        custom_op.parameter_map["auto_tune_mode"].s = tf.compat.as_bytes("RL,GA")
    
    with tf.Session(config=config) as sess:
        if FLAGS.restore:
            print('continue training from previous checkpoint')
            ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_path)
            saver.restore(sess, ckpt)
