Skip to content

Commit

Permalink
Internal change.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 321256761
  • Loading branch information
Lamtharn (Hanoi) Hantrakul authored and Magenta Team committed Jul 14, 2020
1 parent 2de4e1e commit ed3764c
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion ddsp/training/train_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import tensorflow.compat.v2 as tf



# ---------------------- Helper Functions --------------------------------------
def get_strategy(tpu='', gpus=None):
"""Create a distribution strategy.
Expand All @@ -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)
Expand Down

0 comments on commit ed3764c

Please sign in to comment.