Adding TF_USE_LEGACY_KERAS flag to DLRM runs

GoogleCloudPlatform · Feb 10, 2024 · 108ef5e · 108ef5e
1 parent 62a169c
commit 108ef5e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dags/solutions_team/configs/tensorflow/solutionsteam_tf_nightly_supported_config.py b/dags/solutions_team/configs/tensorflow/solutionsteam_tf_nightly_supported_config.py
@@ -279,7 +279,7 @@ def get_tf_dlrm_config(
       "sudo chmod -R 777 /tmp/",
       (
           f"cd /usr/share/tpu/models && {env_variable} &&"
-          " PYTHONPATH='.' python3 official/recommendation/ranking/train.py"
+          " TF_USE_LEGACY_KERAS=1 PYTHONPATH='.' python3 official/recommendation/ranking/train.py"
           f" --tpu={tpu_name} --model_dir=/tmp/output {extraFlags}"
           " --params_override='%s'" % str(params_override)
       ),