diff options
Diffstat (limited to 'tensorflow/python/debug/lib/session_debug_testlib.py')
| -rw-r--r-- | tensorflow/python/debug/lib/session_debug_testlib.py | 3 |
1 files changed, 2 insertions, 1 deletions
diff --git a/tensorflow/python/debug/lib/session_debug_testlib.py b/tensorflow/python/debug/lib/session_debug_testlib.py index 19d4bcae1b..a219ddf9f2 100644 --- a/tensorflow/python/debug/lib/session_debug_testlib.py +++ b/tensorflow/python/debug/lib/session_debug_testlib.py @@ -81,7 +81,8 @@ class SessionDebugTestBase(test_util.TensorFlowTestCase): if test.is_gpu_available(): cls._expected_partition_graph_count = 2 cls._expected_num_devices = 2 - cls._main_device = "/job:localhost/replica:0/task:0/gpu:0" + gpu_name = test_util.gpu_device_name() + cls._main_device = "/job:localhost/replica:0/task:0" + gpu_name else: cls._expected_partition_graph_count = 1 cls._expected_num_devices = 1 |