Update baseimage arg to build time as well

Rohan-Bierneni · Rohan-Bierneni · commit ad0858a07c14 · 2025-12-16T02:06:44.000Z
diff --git a/maxdiffusion_jax_ai_image_tpu.Dockerfile b/maxdiffusion_jax_ai_image_tpu.Dockerfile
@@ -3,6 +3,8 @@ ARG JAX_AI_IMAGE_BASEIMAGE
 # JAX AI Base Image
 FROM $JAX_AI_IMAGE_BASEIMAGE
 
+ARG JAX_AI_IMAGE_BASEIMAGE
+
 ARG COMMIT_HASH
 
 ENV COMMIT_HASH=$COMMIT_HASH
