From 00192f2d02aeb5f9afa6356e6a9f561e2a622e31 Mon Sep 17 00:00:00 2001
From: Dustin H <herbison@google.com>
Date: Tue, 17 Dec 2024 12:18:58 -0500
Subject: [PATCH] Fix CUDA_HOME env var

---
 Dockerfile.tmpl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Dockerfile.tmpl b/Dockerfile.tmpl
index c12a2976..4d1cbc2f 100644
--- a/Dockerfile.tmpl
+++ b/Dockerfile.tmpl
@@ -202,7 +202,7 @@ ENV GIT_COMMIT=${GIT_COMMIT} \
 RUN echo "$GIT_COMMIT" > /etc/git_commit && echo "$BUILD_DATE" > /etc/build_date
 
 {{ if eq .Accelerator "gpu" }}
-    # Add the CUDA home.
-    CUDA_HOME=/usr/local/cuda
+# Add the CUDA home.
+ENV CUDA_HOME=/usr/local/cuda
 {{ end }}
 ENTRYPOINT ["/usr/bin/env"]