mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update the list of default CUDA capabilities used for wheel builds to match build.py.

Browse Source
This commit is contained in:
Peter Hawkins 2022-02-15 09:23:28 -05:00
parent b75b0c04de
commit 2e0cfe8e42
2 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.bazelrc
View File

@ -50,7 +50,7 @@ build:native_arch_posix --host_copt=-march=native
build:mkl_open_source_only --define=tensorflow_mkldnn_contraction_kernel=1
build:cuda --repo_env TF_NEED_CUDA=1
build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="3.5,5.2,6.0,6.1,7.0"
build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="3.5,5.2,6.0,7.0,8.0"
build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
build:cuda --@local_config_cuda//:enable_cuda
build:cuda --define=xla_python_enable_gpu=true

2
build/build.py
View File

@ -375,6 +375,8 @@ def main():
"--cudnn_version",
default=None,
help="CUDNN version, e.g., 8")
# Caution: if changing the default list of CUDA capabilities, you should also
# update the list in .bazelrc, which is used for wheel builds.
parser.add_argument(
"--cuda_compute_capabilities",
default="3.5,5.2,6.0,7.0,8.0",