Add params["tpus"]

undertherain · May 20, 2020 · 5b088b2 · 5b088b2
1 parent 4f8a2e1
commit 5b088b2
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/benchmarker/modules/do_tensorflow.py b/benchmarker/modules/do_tensorflow.py
@@ -29,8 +29,11 @@ def get_strategy(self):
             strategy = tf.distribute.experimental.TPUStrategy(tpu)
             worker_str = tpu.cluster_spec().as_dict()["worker"]
             rep = strategy.num_replicas_in_sync
-            self.params["device"] = "COLAB TPU: {} ".format(worker_str)
-            self.params["device"] += "num_replicas_in_sync: {}".format(rep)
+            self.params["device"] = "COLAB TPU"
+            self.params["tpus"] = {
+                "worker_srt": worker_str,
+                "num_replicas_in_sync": rep,
+            }
         elif len(gpus) > 1:  # multiple GPUs in one VM
             strategy = tf.distribute.MirroredStrategy(gpus)
         else:  # default strategy that works on CPU and single GPU