Avoid loading tf2tensorrt._pywrap_py_utils on windows to workaround

b/173044576.

PiperOrigin-RevId: 342081209
Change-Id: Ifd804f04e6789141587481f62807ed1d6f877b4e
This commit is contained in:
Bixia Zheng 2020-11-12 10:41:54 -08:00 committed by TensorFlower Gardener
parent b081dce0d5
commit c961f9c630

View File

@ -18,7 +18,8 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensorflow.compiler.tf2tensorrt._pywrap_py_utils import is_tensorrt_enabled
import platform
from tensorflow.python.platform import build_info
from tensorflow.python.platform import test
@ -30,8 +31,13 @@ class BuildInfoTest(test.TestCase):
test.is_built_with_rocm())
self.assertEqual(build_info.build_info['is_cuda_build'],
test.is_built_with_cuda())
self.assertEqual(build_info.build_info['is_tensorrt_build'],
is_tensorrt_enabled())
# TODO(b/173044576): make the test work for Windows.
if platform.system() != 'Windows':
# pylint: disable=g-import-not-at-top
from tensorflow.compiler.tf2tensorrt._pywrap_py_utils import is_tensorrt_enabled
self.assertEqual(build_info.build_info['is_tensorrt_build'],
is_tensorrt_enabled())
def testDeterministicOrder(self):
# The dict may contain other keys depending on the platform, but the ones