chore: Update the use of MultiWorkerMirroredStrategy in TF (#1715)

kubeflow · Dec 31, 2022 · 531c19d · 531c19d
1 parent 84961f2
commit 531c19d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/tensorflow/distribution_strategy/keras-API/multi_worker_strategy-with-keras.py b/examples/tensorflow/distribution_strategy/keras-API/multi_worker_strategy-with-keras.py
@@ -73,8 +73,8 @@ def main(args):
   # MultiWorkerMirroredStrategy creates copies of all variables in the model's
   # layers on each device across all workers
   # if your GPUs don't support NCCL, replace "communication" with another
-  strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
-      communication=tf.distribute.experimental.CollectiveCommunication.NCCL)
+  strategy = tf.distribute.MultiWorkerMirroredStrategy(
+      communication_options=tf.distribute.experimental.CommunicationOptions(implementation=tf.distribute.experimental.CollectiveCommunication.AUTO))
 
   BATCH_SIZE_PER_REPLICA = 64
   BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync