fix: setup strategy with set visible devices

TensorSpeech · Jun 15, 2024 · 0e13f76 · 0e13f76
1 parent 31312d5
commit 0e13f76
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensorflow_asr/utils/env_util.py b/tensorflow_asr/utils/env_util.py
@@ -112,7 +112,7 @@ def setup_strategy(
     available_devices = setup_devices(devices)
     if len(available_devices) == 1:
         return tf.distribute.get_strategy()
-    return tf.distribute.MirroredStrategy(devices=[d.name for d in available_devices])
+    return tf.distribute.MultiWorkerMirroredStrategy()
 
 
 def has_devices(