From 8c44532ce09e4c369a27120059f86fe1bcbee728 Mon Sep 17 00:00:00 2001 From: Param Bole Date: Thu, 9 Jan 2025 21:27:25 +0000 Subject: [PATCH] Updating Image tags --- dags/common/vm_resource.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dags/common/vm_resource.py b/dags/common/vm_resource.py index e8d6ff5e..2a61f275 100644 --- a/dags/common/vm_resource.py +++ b/dags/common/vm_resource.py @@ -297,7 +297,7 @@ class DockerImage(enum.Enum): f"xla:nightly_3.10_tpuvm_{datetime.datetime.today().strftime('%Y%m%d')}" ) AXLEARN_TPU_JAX_STABLE_STACK = ( - "us-docker.pkg.dev/tpu-prod-env-multipod/bite/tpu/jax0.4.35-rev1:" + "us-docker.pkg.dev/tpu-prod-env-multipod/bite/tpu/axlearn:" f"{datetime.datetime.today().strftime('%Y-%m-%d')}" ) AXLEARN_GPU_JAX_NIGHTLY = ( @@ -313,7 +313,7 @@ class DockerImage(enum.Enum): f"{datetime.datetime.today().strftime('%Y-%m-%d')}" ) MAXDIFFUSION_TPU_JAX_STABLE_STACK = ( - "gcr.io/tpu-prod-env-multipod/maxdiffusion_jax_stable_stack_0.4.35:" + "gcr.io/tpu-prod-env-multipod/maxdiffusion_jax_stable_stack:" f"{datetime.datetime.today().strftime('%Y-%m-%d')}" ) MAXDIFFUSION_TPU_JAX_NIGHTLY = (