def maybe_cloud_tpu(): """If FLAGS.cloud_tpu is set, setup Cloud instances.""" if not FLAGS.cloud_tpu: yield return tf.logging.info("Running on Cloud TPU") if (not FLAGS.data_dir.startswith("gs://") or not FLAGS.output_dir.startswith("gs://")): raise ValueError("To run on Cloud TPUs, data_dir and output_dir need to " "be gs:// paths, i.e. on Google Cloud Storage.") FLAGS.use_tpu = True with cloud_tpu.cloud_tpu( FLAGS.cloud_vm_name, FLAGS.cloud_tpu_name, delete_on_done=FLAGS.cloud_delete_on_done) as tpu_master: FLAGS.master = tpu_master yield