diff options
Diffstat (limited to 'tensorflow/contrib/BUILD')
| -rw-r--r-- | tensorflow/contrib/BUILD | 5 |
1 files changed, 4 insertions, 1 deletions
diff --git a/tensorflow/contrib/BUILD b/tensorflow/contrib/BUILD index 6b3343bb2f..bab37e8906 100644 --- a/tensorflow/contrib/BUILD +++ b/tensorflow/contrib/BUILD @@ -7,6 +7,7 @@ package(default_visibility = ["//tensorflow:__subpackages__"]) load("//third_party/mpi:mpi.bzl", "if_mpi") load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda") +load("@local_config_tensorrt//:build_defs.bzl", "if_tensorrt") py_library( name = "contrib_py", @@ -107,7 +108,9 @@ py_library( "//tensorflow/contrib/training:training_py", "//tensorflow/contrib/util:util_py", "//tensorflow/python:util", - ] + if_mpi(["//tensorflow/contrib/mpi_collectives:mpi_collectives_py"]), + ] + if_mpi(["//tensorflow/contrib/mpi_collectives:mpi_collectives_py"]) + if_tensorrt([ + "//tensorflow/contrib/tensorrt:init_py", + ]), ) cc_library( |