Commit 2b79cb39 authored by Amir MOHAMMADI's avatar Amir MOHAMMADI
Browse files

explicitly use nccl

parent f1fa89e9
import logging
import tensorflow as tf
logger = logging.getLogger(__name__)
def strategy_fn():
print("creating strategy")
strategy = tf.distribute.MultiWorkerMirroredStrategy()
print("strategy created")
print("Creating MultiWorkerMirroredStrategy strategy.")
strategy = tf.distribute.MultiWorkerMirroredStrategy(
communication_options=tf.distribute.experimental.CommunicationOptions(
implementation=tf.distribute.experimental.CollectiveCommunication.NCCL
)
)
print("MultiWorkerMirroredStrategy strategy created.")
return strategy
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment