don't set the tpu_job_name for pre-training

google-research · Mar 11, 2020 · 4e0f2a5 · 4e0f2a5
1 parent cbf6f73
commit 4e0f2a5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/configure_pretraining.py b/configure_pretraining.py
@@ -81,6 +81,7 @@ def __init__(self, model_name, data_dir, **kwargs):
     # TPU settings
     self.use_tpu = False
     self.num_tpu_cores = 1
+    self.tpu_job_name = None
     self.tpu_name = None  # cloud TPU to use for training
     self.tpu_zone = None  # GCE zone where the Cloud TPU is located in
     self.gcp_project = None  # project name for the Cloud TPU-enabled project

diff --git a/run_pretraining.py b/run_pretraining.py
@@ -323,8 +323,7 @@ def train_or_eval(config: configure_pretraining.PretrainingConfig):
       iterations_per_loop=config.iterations_per_loop,
       num_shards=(config.num_tpu_cores if config.do_train else
                   config.num_tpu_cores),
-      tpu_job_name=("train_tpu_worker" if config.do_train else
-                    "lm_eval_tpu_worker"),
+      tpu_job_name=config.tpu_job_name,
       per_host_input_for_training=is_per_host)
   run_config = tf.estimator.tpu.RunConfig(
       cluster=tpu_cluster_resolver,