Call connect_to_cluster for initialization test.

PiperOrigin-RevId: 298442916
Change-Id: I588c258768cda131844f26ff2f2abf2de3a97119
This commit is contained in:
Ken Franko 2020-03-02 14:10:56 -08:00 committed by TensorFlower Gardener
parent ef6ef73686
commit 5c1f1be66d

View File

@ -19,6 +19,7 @@ from __future__ import division
from __future__ import print_function
from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
from tensorflow.python.eager import remote
from tensorflow.python.eager import test
from tensorflow.python.platform import flags
from tensorflow.python.platform import tf_logging as logging
@ -39,6 +40,7 @@ class TpuStrategyTest(test.TestCase):
zone=FLAGS.zone,
project=FLAGS.project,
)
remote.connect_to_cluster(resolver)
tpu_strategy_util.initialize_tpu_system(resolver)
with test.mock.patch.object(logging, "warning") as mock_log: