Consider installing local CUDA variant when building GPU image

[chajath]

We currently install jax[cuda12_pip] in setup.sh. However, since the base image comes with its own CUDA installation, this creates a version conflict, generating warnings such as:

2024-03-11 22:25:16.557390: I external/xla/xla/stream_executor/cuda/cuda_dnn.cc:517] Loaded cuDNN version 8907
2024-03-11 22:25:16.563077: W external/xla/xla/service/gpu/nvptx_compiler.cc:742] The NVIDIA driver's CUDA version is 12.2 which is older than the ptxas CUDA version (12.4.99). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.

Installing jax[cuda12_local] seems to mitigate this problem, although we should be careful in pinning aqtp version when doing so (see #500)