fix nightly

GoogleCloudPlatform · Nov 26, 2024 · 39c9536 · 39c9536
1 parent fbb1de9
commit 39c9536
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/dags/solutions_team/solutionsteam_tf_nightly_supported.py b/dags/solutions_team/solutionsteam_tf_nightly_supported.py
@@ -41,20 +41,23 @@
       tpu_zone=Zone.US_CENTRAL1_C.value,
       time_out_in_min=60,
       global_batch_size=1024,
+      runtime_version=RuntimeVersion.V2_ALPHA_TPUV5.value,
   )
 
   tf_resnet_v3_8 = tf_config.get_tf_resnet_config(
       tpu_version=TpuVersion.V3,
       tpu_cores=8,
       tpu_zone=Zone.US_EAST1_D.value,
       time_out_in_min=60,
+      runtime_version=RuntimeVersion.V2_ALPHA_TPUV5.value,
   )
 
   tf_resnet_v4_8 = tf_config.get_tf_resnet_config(
       tpu_version=TpuVersion.V4,
       tpu_cores=8,
       tpu_zone=Zone.US_CENTRAL2_B.value,
       time_out_in_min=60,
+      runtime_version=RuntimeVersion.V2_ALPHA_TPUV5.value,
   )
 
   tf_resnet_v4_32 = tf_config.get_tf_resnet_config(