diff options
Diffstat (limited to 'tensorflow/python/profiler/model_analyzer_test.py')
| -rw-r--r-- | tensorflow/python/profiler/model_analyzer_test.py | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/tensorflow/python/profiler/model_analyzer_test.py b/tensorflow/python/profiler/model_analyzer_test.py index f9891f3b1e..c0e16ca536 100644 --- a/tensorflow/python/profiler/model_analyzer_test.py +++ b/tensorflow/python/profiler/model_analyzer_test.py @@ -106,7 +106,7 @@ class PrintModelAnalysisTest(test.TestCase): # Make sure time is profiled. gap = 1 if test.is_gpu_available() else 2 for i in range(3, 6, gap): - mat = re.search('(.*)[um]s/(.*)[um]s', metrics[i]) + mat = re.search('(.*)(?:us|ms|sec)/(.*)(?:us|ms|sec)', metrics[i]) self.assertGreater(float(mat.group(1)), 0.0) self.assertGreater(float(mat.group(2)), 0.0) # Make sure device is profiled. |