From ed3764c50da66d71e38d772ec71b1678c442c58d Mon Sep 17 00:00:00 2001 From: "Lamtharn (Hanoi) Hantrakul" Date: Tue, 14 Jul 2020 16:18:49 -0700 Subject: [PATCH] Internal change. PiperOrigin-RevId: 321256761 --- ddsp/training/train_util.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ddsp/training/train_util.py b/ddsp/training/train_util.py index a7eb17e5..b6f5b046 100644 --- a/ddsp/training/train_util.py +++ b/ddsp/training/train_util.py @@ -23,6 +23,7 @@ import tensorflow.compat.v2 as tf + # ---------------------- Helper Functions -------------------------------------- def get_strategy(tpu='', gpus=None): """Create a distribution strategy. @@ -40,7 +41,7 @@ def get_strategy(tpu='', gpus=None): resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=tpu) tf.config.experimental_connect_to_cluster(resolver) tf.tpu.experimental.initialize_tpu_system(resolver) - strategy = tf.distribute.experimental.TPUStrategy(resolver) + strategy = tf.distribute.TPUStrategy(resolver) elif gpus: for gpu_address in gpus: logging.info('Use GPU at %s', gpu_address)