From 7d82dbb42744a21ff05924e973e57a68465f3347 Mon Sep 17 00:00:00 2001 From: Igor Saprykin Date: Fri, 1 Dec 2017 19:44:35 -0800 Subject: Fix a replicate_model_fn_test that is dependent on the number of hardware GPUs. It has been causing failures when run on a machine with 4 GPUs. PiperOrigin-RevId: 177670759 --- .../contrib/estimator/python/estimator/replicate_model_fn_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py index 662021853d..91e4b9ba7d 100644 --- a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py +++ b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py @@ -288,7 +288,7 @@ class ReplicateModelTest(test_util.TensorFlowTestCase): with self.test_session() as session: replicated_model_fn = replicate_model_fn.replicate_model_fn( - self.model_fn, self.optimizer_fn) + self.model_fn, self.optimizer_fn, devices=['/gpu:0']) estimator_spec = replicated_model_fn( features, labels, model_fn_lib.ModeKeys.TRAIN, self.params) session.run(variables.global_variables_initializer()) -- cgit v1.2.3