Use context.config instead of a clean tf.ConfigProto in test_util
PiperOrigin-RevId: 251371530
This commit is contained in:
parent
595cb7ecd9
commit
2892a6844c
@ -49,7 +49,6 @@ from google.protobuf import descriptor_pool
|
||||
from google.protobuf import text_format
|
||||
|
||||
from tensorflow.core.framework import graph_pb2
|
||||
from tensorflow.core.protobuf import config_pb2
|
||||
from tensorflow.core.protobuf import rewriter_config_pb2
|
||||
from tensorflow.python import pywrap_tensorflow
|
||||
from tensorflow.python import tf2
|
||||
@ -2741,12 +2740,10 @@ class TensorFlowTestCase(googletest.TestCase):
|
||||
# will be used even when a specific device is supposed to be used.
|
||||
allow_soft_placement = not force_gpu
|
||||
if config is None:
|
||||
config = config_pb2.ConfigProto()
|
||||
config = context.context().config
|
||||
config.allow_soft_placement = allow_soft_placement
|
||||
config.gpu_options.per_process_gpu_memory_fraction = 0.3
|
||||
elif not allow_soft_placement and config.allow_soft_placement:
|
||||
config_copy = config_pb2.ConfigProto()
|
||||
config_copy.CopyFrom(config)
|
||||
config_copy = context.context().config
|
||||
config = config_copy
|
||||
config.allow_soft_placement = False
|
||||
# Don't perform optimizations for tests so we don't inadvertently run
|
||||
|
Loading…
x
Reference in New Issue
Block a user