From 915bef5c6a0c3c56a8647d84be19ae2def6369dd Mon Sep 17 00:00:00 2001 From: Rick Chao Date: Wed, 20 Nov 2019 23:39:46 -0800 Subject: [PATCH] In NamedGPUCombination.should_execute_combination, check if `number_of_required_gpus > 0` before making a call to context.num_gpus(), since if `number_of_required_gpus` is zero, it does not even need to query the number of gpus available. This alleviates the requirement that collective ops needs to be configured at program startup for cases where gpu is not needed, because this code path is run when using combinations, which forces collective ops to be configured even before the test body is run. PiperOrigin-RevId: 281685344 Change-Id: Ic53c3a50c95c512bd88cf685a6f50c1867b469b7 --- tensorflow/python/distribute/combinations.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/distribute/combinations.py b/tensorflow/python/distribute/combinations.py index 0c1fdbe82fd..c4915eb9d80 100644 --- a/tensorflow/python/distribute/combinations.py +++ b/tensorflow/python/distribute/combinations.py @@ -79,7 +79,8 @@ class NamedGPUCombination(test_combinations.TestCombination): if not number_of_required_gpus and GPUCombination.GPU_TEST: return (False, "Test that doesn't require GPUs.") - elif context.num_gpus() < number_of_required_gpus: + elif (number_of_required_gpus > 0 + and context.num_gpus() < number_of_required_gpus): return (False, ("Only {} of {} required GPUs are available.".format( context.num_gpus(), number_of_required_gpus))) else: