diff options
Diffstat (limited to 'tensorflow/tensorflow.bzl')
| -rw-r--r-- | tensorflow/tensorflow.bzl | 2 |
1 files changed, 2 insertions, 0 deletions
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl index c3bc9ccd45..6bb393a3f4 100644 --- a/tensorflow/tensorflow.bzl +++ b/tensorflow/tensorflow.bzl @@ -922,6 +922,7 @@ def tf_gpu_kernel_library(srcs, hdrs=[], **kwargs): copts = copts + _cuda_copts() + if_cuda(cuda_copts) + tf_copts() + kwargs["features"] = kwargs.get("features", []) + ["-use_header_modules"] native.cc_library( srcs=srcs, @@ -1305,6 +1306,7 @@ def tf_custom_op_library(name, srcs=[], gpu_srcs=[], deps=[], linkopts=[]): name=basename + "_gpu", srcs=gpu_srcs, copts=_cuda_copts() + if_tensorrt(["-DGOOGLE_TENSORRT=1"]), + features = if_cuda(["-use_header_modules"]), deps=deps + if_cuda(cuda_deps)) cuda_deps.extend([":" + basename + "_gpu"]) |