experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix a multi-gpu test failure.

Browse Source 
The test uses tf.constant as input to all_reduce in pure eager mode, however in eager mode tf.constant always creates host tensors regardless of the enclosing device scope. This leads to NCCL error.

PiperOrigin-RevId: 357252841
Change-Id: Iddaf5f52fe6634ec29dd385a9fa034761f3df91f
This commit is contained in:
Chenkai Kuang 2021-02-12 13:06:24 -08:00 committed by TensorFlower Gardener
parent 8e4ba814c2
commit aa9bd19fe3
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/python/distribute/strategy_common_test.py
View File

@ -131,7 +131,7 @@ class ReplicaCtxAllReduceTest(test.TestCase, parameterized.TestCase):
def fn():
def replica_fn():
value = constant_op.constant(1.0)
value = array_ops.identity(1.0)
reduced = strategy.extended._replica_ctx_all_reduce('SUM', value)
return reduced
@ -152,7 +152,7 @@ class ReplicaCtxAllReduceTest(test.TestCase, parameterized.TestCase):
def fn():
def replica_fn():
value = (constant_op.constant(1.0), constant_op.constant(2.0))
value = (array_ops.identity(1.0), array_ops.identity(1.0))
reduced = strategy.extended._replica_ctx_all_reduce('SUM', value)
return reduced