Merge pull request #22136 from jackonan:master

PiperOrigin-RevId: 222866762

tensorflow/python/ops/resources.py

@@ -21,6 +21,7 @@ from __future__ import division
 from __future__ import print_function
 
 import collections
+import os
 
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
@@ -86,7 +87,9 @@ def report_uninitialized_resources(resource_list=None,
     resource_list = shared_resources() + local_resources()
   with ops.name_scope(name):
     # Run all operations on CPU
-    with ops.device("/cpu:0"):
+    local_device = os.environ.get(
+        "TF_DEVICE_FOR_UNINITIALIZED_VARIABLE_REPORTING", "/cpu:0")
+    with ops.device(local_device):
       if not resource_list:
         # Return an empty tensor so we only need to check for returned tensor
         # size being 0 as an indication of model ready.

tensorflow/python/ops/variables.py

@@ -19,6 +19,7 @@ from __future__ import print_function
 
 import enum  # pylint: disable=g-bad-import-order
 
+import os
 import six
 
 from tensorflow.core.framework import attr_value_pb2
@@ -2963,7 +2964,9 @@ def report_uninitialized_variables(var_list=None,
     # Run all operations on CPU
     if var_list:
       init_vars = [state_ops.is_variable_initialized(v) for v in var_list]
-    with ops.device("/cpu:0"):
+    local_device = os.environ.get(
+        "TF_DEVICE_FOR_UNINITIALIZED_VARIABLE_REPORTING", "/cpu:0")
+    with ops.device(local_device):
       if not var_list:
         # Return an empty tensor so we only need to check for returned tensor
         # size being 0 as an indication of model ready.