summaryrefslogtreecommitdiff
path: root/var
diff options
context:
space:
mode:
Diffstat (limited to 'var')
-rw-r--r--var/spack/repos/builtin/packages/py-tensorflow/package.py1
1 files changed, 1 insertions, 0 deletions
diff --git a/var/spack/repos/builtin/packages/py-tensorflow/package.py b/var/spack/repos/builtin/packages/py-tensorflow/package.py
index fd152cb89b..1d86e81fd0 100644
--- a/var/spack/repos/builtin/packages/py-tensorflow/package.py
+++ b/var/spack/repos/builtin/packages/py-tensorflow/package.py
@@ -562,6 +562,7 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
for pkg_dep in rocm_dependencies:
pkg_dep_cap = pkg_dep.upper().replace("-", "_")
env.set(f"{pkg_dep_cap}_PATH", spec[pkg_dep].prefix)
+ env.set("TF_ROCM_AMDGPU_TARGETS", ",".join(self.spec.variants["amdgpu_target"].value))
else:
env.set("TF_NEED_ROCM", "0")