Fix numba-cuda CPU import crash: install only on GPU image

numba-cuda v0.30.0 depends on cuda-bindings which requires libcudart.so
at import time. This crashes on the CPU image where no CUDA runtime is
installed. Move numba-cuda install into the GPU-only section while
keeping the numba upgrade for both CPU and GPU.

b/485275559

Author: stevemessick

Dockerfile.tmpl

@@ -28,8 +28,8 @@ RUN uv pip install --no-build-isolation --no-cache --system "git+https://github.
 # b/468367647: Unpin protobuf, version greater than v5.29.5 causes issues with numerous packages
 RUN uv pip install --system --force-reinstall --no-cache --no-deps torchtune
 RUN uv pip install --system --force-reinstall --no-cache "protobuf==5.29.5"
-# b/493600019: Colab base image ships numba/numba-cuda that do not support NumPy 2.4; upgrade both.
-RUN uv pip install --system --force-reinstall --no-cache numba numba-cuda
+# b/493600019: Colab base image ships numba that does not support NumPy 2.4; upgrade to latest.
+RUN uv pip install --system --force-reinstall --no-cache numba
 
 # Adding non-package dependencies:
 ADD clean-layer.sh  /tmp/clean-layer.sh
@@ -40,6 +40,8 @@ ARG PACKAGE_PATH=/usr/local/lib/python3.12/dist-packages
 
 # Install GPU-specific non-pip packages.
 {{ if eq .Accelerator "gpu" }}
+# b/493600019: numba-cuda v0.30.0 fixes np.trapz removal in NumPy 2.4 but requires libcudart.so (GPU only).
+RUN uv pip install --system --force-reinstall --no-cache numba-cuda
 RUN uv pip install --system --no-cache "pycuda"
 {{ end }}