diff options
Diffstat (limited to 'tensorflow/contrib/tensorrt/BUILD')
| -rw-r--r-- | tensorflow/contrib/tensorrt/BUILD | 31 |
1 files changed, 31 insertions, 0 deletions
diff --git a/tensorflow/contrib/tensorrt/BUILD b/tensorflow/contrib/tensorrt/BUILD index 122a67a407..9e8979bce4 100644 --- a/tensorflow/contrib/tensorrt/BUILD +++ b/tensorflow/contrib/tensorrt/BUILD @@ -19,6 +19,7 @@ load( "tf_gen_op_libs", "tf_gen_op_wrapper_py", ) +load("//tensorflow:tensorflow.bzl", "cuda_py_test") load("//tensorflow:tensorflow.bzl", "cuda_py_tests") load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test") load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") @@ -181,7 +182,12 @@ py_library( srcs_version = "PY2AND3", deps = [ ":wrap_conversion", + "//tensorflow/python:graph_util", + "//tensorflow/python:session", "//tensorflow/python:tf_optimizer", + "//tensorflow/python/saved_model:builder", + "//tensorflow/python/saved_model:loader", + "//tensorflow/python/saved_model:tag_constants", ], ) @@ -410,6 +416,31 @@ py_library( ], ) +cuda_py_test( + name = "trt_convert_test", + srcs = ["python/trt_convert_test.py"], + additional_deps = [ + ":trt_convert_py", + ":trt_ops_py", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:graph_util", + "//tensorflow/python/saved_model:builder", + "//tensorflow/python/saved_model:loader", + "//tensorflow/python/saved_model:signature_constants", + "//tensorflow/python/saved_model:signature_def_utils", + "//tensorflow/python/saved_model:tag_constants", + "//tensorflow/python/saved_model:utils", + "//tensorflow/python/tools:freeze_graph_lib", + "//tensorflow/python/tools:saved_model_utils", + ], + tags = [ + "no_cuda_on_cpu_tap", + "no_windows", + "nomac", + ], +) + cuda_py_tests( name = "tf_trt_integration_test", srcs = [ |