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

[TF:TRT] Fix a python test.

Browse Source 
For TF2, TRTEngineOp nodes should only exist inside functions for
signatures. Avoid looking for TRTEngineOp nodes outside functions.

PiperOrigin-RevId: 358194388
Change-Id: I123b6e366a130c87b1e0f5bc64336bf21f715b20
This commit is contained in:
Bixia Zheng 2021-02-18 09:21:13 -08:00 committed by TensorFlower Gardener
parent 55367cfc7d
commit c92320d7a4
1 changed files with 1 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/python/compiler/tensorrt/test/tf_trt_integration_test_base.py
View File

@ -865,9 +865,7 @@ class TfTrtIntegrationTestBase(test_util.TensorFlowTestCase):
return
expected_engines = self.ExpectedEnginesToBuild(run_params)
all_op_names = [node.name for node in gdef_to_verify.node]
trt_op_names = [
node.name for node in gdef_to_verify.node if node.op == "TRTEngineOp"
]
trt_op_names = []
for func in gdef_to_verify.library.function:
if not re.search(r"TRTEngineOp_\d+_\d+_native_segment",
func.signature.name):