Reorder FROM & ENV so that build works correctly

Author: djherbis

Dockerfile.tmpl

@@ -8,13 +8,18 @@ ARG TORCHAUDIO_VERSION
 ARG TORCHTEXT_VERSION
 ARG TORCHVISION_VERSION
 
-# Ensures shared libraries installed with conda can be found by the dynamic link loader.
-ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/opt/conda/lib"
-
 {{ if eq .Accelerator "gpu" }}
 FROM gcr.io/kaggle-images/python-lightgbm-whl:${GPU_BASE_IMAGE_NAME}-${BASE_IMAGE_TAG}-${LIGHTGBM_VERSION} AS lightgbm_whl
 FROM gcr.io/kaggle-images/python-torch-whl:${GPU_BASE_IMAGE_NAME}-${BASE_IMAGE_TAG}-${TORCH_VERSION} AS torch_whl
 FROM ${BASE_IMAGE_REPO}/${GPU_BASE_IMAGE_NAME}:${BASE_IMAGE_TAG}
+{{ else }}
+FROM ${BASE_IMAGE_REPO}/${CPU_BASE_IMAGE_NAME}:${BASE_IMAGE_TAG}
+{{ end }}
+
+# Ensures shared libraries installed with conda can be found by the dynamic link loader.
+ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/opt/conda/lib"
+
+{{ if eq .Accelerator "gpu" }}
 ARG CUDA_MAJOR_VERSION
 ARG CUDA_MINOR_VERSION
 ENV CUDA_MAJOR_VERSION=${CUDA_MAJOR_VERSION}
@@ -25,12 +30,10 @@ ENV PATH=/opt/bin:${PATH}
 ENV LD_LIBRARY_PATH_NO_STUBS="$LD_LIBRARY_PATH"
 ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/lib64/stubs"
 RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1
-{{ else }}
-FROM ${BASE_IMAGE_REPO}/${CPU_BASE_IMAGE_NAME}:${BASE_IMAGE_TAG}
 {{ end }}
+
 # Keep these variables in sync if base image is updated.
-# TODO(rosbo): Still 2.9.0 rc2. Wait for the 2.9.1
-ENV TENSORFLOW_VERSION=2.9.0
+ENV TENSORFLOW_VERSION=2.9.1
 
 # We need to redefine the ARG here to get the ARG value defined above the FROM instruction.
 # See: https://docs.docker.com/engine/reference/builder/#understand-how-arg-and-from-interact