Avoid loading tf2tensorrt._pywrap_py_utils on windows to workaround
b/173044576. PiperOrigin-RevId: 342081209 Change-Id: Ifd804f04e6789141587481f62807ed1d6f877b4e
This commit is contained in:
parent
b081dce0d5
commit
c961f9c630
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user