correct flag name category

mlcommons · Dec 20, 2023 · fcd172c · fcd172c
1 parent cd6d7cb
commit fcd172c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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"