From fcd172cae6a57bf11cc3e357658688556ff29ec3 Mon Sep 17 00:00:00 2001
From: Priya Kasimbeg <kasimbeg@google.com>
Date: Wed, 20 Dec 2023 01:26:19 +0000
Subject: [PATCH] correct flag name category

---
 submission_runner.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/submission_runner.py b/submission_runner.py
index d32a7824f..eb89e9cf5 100644
--- a/submission_runner.py
+++ b/submission_runner.py
@@ -54,7 +54,7 @@
 os.environ['TF_XLA_FLAGS'] = '--xla_gpu_enable_triton_gemm=false'
 
 # Flags to optimize JAX GPU
-os.environ["XLA_FLAGS"] = (
+os.environ["TF_XLA_FLAGS"] = (
   "--xla_gpu_enable_async_collectives=true"
   " --xla_gpu_enable_latency_hiding_scheduler=true"
   " --xla_gpu_enable_highest_priority_async_stream=true"