Skip to content
Snippets Groups Projects
Commit 11d3abe8 authored by Eric Müller's avatar Eric Müller :mountain_bicyclist:
Browse files

fix(py-tensorflow): missing conf env vars for @2.18:

parent added671
No related branches found
No related tags found
1 merge request!655Draft: Fix py-{jax{,lib},tensorflow}+cuda
...@@ -728,6 +728,18 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension): ...@@ -728,6 +728,18 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
# Please specify the cuDNN version you want to use # Please specify the cuDNN version you want to use
env.set("TF_CUDNN_VERSION", spec["cudnn"].version.up_to(1)) env.set("TF_CUDNN_VERSION", spec["cudnn"].version.up_to(1))
# Please specify the hermetic CUDA version you want to use
env.set("HERMETIC_CUDA_VERSION", spec["cuda"].version.up_to(3))
# Please specify the hermetic cuDNN version you want to use
env.set("HERMETIC_CUDNN_VERSION", spec["cudnn"].version.up_to(3))
# Please specify the local CUDA path you want to use
env.set("LOCAL_CUDA_PATH", spec["cuda"].prefix)
# Please specify the local CUDNN path you want to use
env.set("LOCAL_CUDNN_PATH", spec["cudnn"].prefix)
if "+nccl" in spec: if "+nccl" in spec:
cuda_paths.append(spec["nccl"].prefix) cuda_paths.append(spec["nccl"].prefix)
...@@ -737,6 +749,9 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension): ...@@ -737,6 +749,9 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
# Please specify the location where NCCL is installed # Please specify the location where NCCL is installed
env.set("NCCL_INSTALL_PATH", spec["nccl"].prefix) env.set("NCCL_INSTALL_PATH", spec["nccl"].prefix)
env.set("NCCL_HDR_PATH", spec["nccl"].prefix.include) env.set("NCCL_HDR_PATH", spec["nccl"].prefix.include)
# Please specify the local NCCL path you want to use
env.set("LOCAL_NCCL_PATH", spec["nccl"].prefix)
else: else:
env.unset("TF_NCCL_VERSION") env.unset("TF_NCCL_VERSION")
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment