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

Merge pull request #44803 from cloudhan:jax_winbuild_cuda

Browse Source 
PiperOrigin-RevId: 342704153
Change-Id: I5fe5dc540df8f83aaa801b6dc9c8591d050233b0
This commit is contained in:
TensorFlower Gardener 2020-11-16 13:23:33 -08:00
commit c481421627
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/compiler/xla/pjrt/BUILD
View File

@ -1,6 +1,6 @@
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "if_nccl")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
package(
default_visibility = ["//tensorflow:internal"],
@ -240,7 +240,7 @@ cc_library(
name = "nvidia_gpu_device",
srcs = ["nvidia_gpu_device.cc"],
hdrs = ["nvidia_gpu_device.h"],
copts = if_cuda(["-DNCCL_ENABLED=1"]),
copts = if_nccl(["-DNCCL_ENABLED=1"]),
deps = [
":pjrt_client",
"//tensorflow/compiler/xla/service/gpu:gpu_executable_run_options",
@ -253,7 +253,7 @@ cc_library(
"//tensorflow/core/common_runtime/gpu:gpu_mem_allocator",
"//tensorflow/core:lib_internal",
"//tensorflow/stream_executor:tf_allocator_adapter",
] + if_cuda(["@local_config_nccl//:nccl"]),
] + if_nccl(["@local_config_nccl//:nccl"]),
)
tf_cc_test(