Skip to content
Snippets Groups Projects
Commit 2b79cb39 authored by Amir MOHAMMADI's avatar Amir MOHAMMADI
Browse files

explicitly use nccl

parent f1fa89e9
Branches
Tags
No related merge requests found
import logging
import tensorflow as tf import tensorflow as tf
logger = logging.getLogger(__name__)
def strategy_fn(): def strategy_fn():
print("creating strategy") print("Creating MultiWorkerMirroredStrategy strategy.")
strategy = tf.distribute.MultiWorkerMirroredStrategy() strategy = tf.distribute.MultiWorkerMirroredStrategy(
print("strategy created") communication_options=tf.distribute.experimental.CommunicationOptions(
implementation=tf.distribute.experimental.CollectiveCommunication.NCCL
)
)
print("MultiWorkerMirroredStrategy strategy created.")
return strategy return strategy
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment