experiments
/
STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Treat "test_xla_gpu" as GPU_TEST in "NamedGPUCombination". 

PiperOrigin-RevId: 309276564
Change-Id: Iea914898c6a090c8397d0ad0890e239948eec943
This commit is contained in:
Chenkai Kuang 2020-04-30 12:25:37 -07:00 committed by TensorFlower Gardener
parent f69af82c3d
commit dfe0f0aa36
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/python/distribute/combinations.py
View File

@ -22,6 +22,7 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import re
import sys
import types
import unittest
@ -94,10 +95,10 @@ class NamedGPUCombination(combinations_lib.TestCombination):
Attributes:
GPU_TEST: The environment is considered to have GPU hardware available if
the name of the program contains "test_gpu".
the name of the program contains "test_gpu" or "test_xla_gpu".
"""
GPU_TEST = "test_gpu" in sys.argv[0]
GPU_TEST = re.search(r"(test_gpu|test_xla_gpu)$", sys.argv[0])
def should_execute_combination(self, kwargs):
distributions = [