diff --git a/packages/py-tensorflow/package.py b/packages/py-tensorflow/package.py index 436a1ea575bbbcb26571198ec178890184a08b13..6b01431b6faa4ce83ebfa02435a54be0e91e5e37 100644 --- a/packages/py-tensorflow/package.py +++ b/packages/py-tensorflow/package.py @@ -728,6 +728,18 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension): # Please specify the cuDNN version you want to use 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: cuda_paths.append(spec["nccl"].prefix) @@ -737,6 +749,9 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension): # Please specify the location where NCCL is installed env.set("NCCL_INSTALL_PATH", spec["nccl"].prefix) 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: env.unset("TF_NCCL_VERSION")