From e576acf5dbd7b800d3b6aa4de4b69952a9e2c0fb Mon Sep 17 00:00:00 2001 From: Brennan Saeta Date: Thu, 6 Dec 2018 10:28:10 -0800 Subject: [PATCH] Internal-only change. PiperOrigin-RevId: 224362520 --- .../distribute/cluster_resolver/tpu_cluster_resolver.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver.py b/tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver.py index fe8e09f5612..1fd674c8a46 100644 --- a/tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver.py +++ b/tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver.py @@ -197,13 +197,14 @@ class TPUClusterResolver(ClusterResolver): elif tpu == 'local' or not tpu: # Google environment, where the TPU is attached to the host. self._environment = 'google' - elif tpu.startswith('/bns'): + elif tpu.startswith('/bns') or tpu.startswith('uptc://'): # Google environment, where we reach the TPU through BNS. self._environment = 'google' # If TPU is in the Google environment or exists locally, we don't use any # RPC layer. - if tpu.startswith('/bns') or tpu == 'local' or not tpu: + if tpu.startswith('/bns') or tpu.startswith( + 'uptc://') or tpu == 'local' or not tpu: self.rpc_layer = None else: self.rpc_layer = 'grpc'