Fix transformer_main.py.

An earlier change made it so running this file would cause an error. The unit tests still passed, as the unit tests do not directly call the main() function PiperOrigin-RevId: 264268400
Navneeth08k · Aug 19, 2019 · b1188d0 · b1188d0
1 parent 67420e1
commit b1188d0
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/official/transformer/v2/transformer_main.py b/official/transformer/v2/transformer_main.py
@@ -403,8 +403,7 @@ def _run_task(task):
     if not flags_obj.distribution_strategy != "tpu":
       _run_task(task)
     else:
-      primary_cpu_task = ("/job:worker"
-                          if flags_obj.use_tpu_2vm_config is not None else "")
+      primary_cpu_task = "/job:worker" if flags_obj.use_tpu_2vm_config else ""
       with tf.device(primary_cpu_task):
         _run_task(task)