From ffdd38af247c8cbfa21f0b4eb0923994654eee9f Mon Sep 17 00:00:00 2001
From: JackCaoG <jackcao@google.com>
Date: Fri, 16 Aug 2024 00:27:31 +0000
Subject: [PATCH] disable test_traceable_collectives

---
 test/run_tests.sh     | 3 ++-
 test/tpu/run_tests.sh | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/test/run_tests.sh b/test/run_tests.sh
index b4cda7c3cef4..51eecb539119 100755
--- a/test/run_tests.sh
+++ b/test/run_tests.sh
@@ -299,7 +299,8 @@ function run_mp_op_tests {
   run_test "$CDIR/test_mp_save.py"
   run_test "$CDIR/test_mp_mesh_reduce.py"
   run_test "$CDIR/test_mp_sync_batch_norm.py"
-  run_test "$CDIR/dynamo/test_traceable_collectives.py"
+  # TODO(JackCaoG): enable this
+  # run_test "$CDIR/dynamo/test_traceable_collectives.py"
   run_test "$CDIR/test_fsdp_auto_wrap.py"
   # run_torchrun "$CDIR/test_mp_early_exit.py"
   run_pt_xla_debug "$CDIR/debug_tool/test_mp_pt_xla_debug.py"
diff --git a/test/tpu/run_tests.sh b/test/tpu/run_tests.sh
index f143ba61353d..56e80689aa8a 100755
--- a/test/tpu/run_tests.sh
+++ b/test/tpu/run_tests.sh
@@ -44,7 +44,8 @@ python3 examples/train_resnet_amp.py
 # Egaer tests will take more HBM, only run them on TPU v4 CI
 TPU_VERSION=$(python -c "import sys; sys.path.remove(''); import torch_xla; print(torch_xla._internal.tpu.version())")
 if [[ -n "$TPU_VERSION" && "$TPU_VERSION" == "4" ]]; then
-    python3 test/dynamo/test_traceable_collectives.py
+    # TODO(JackCaoG): fix this
+    # python3 test/dynamo/test_traceable_collectives.py
     python3 examples/data_parallel/train_resnet_xla_ddp.py
     python3 examples/fsdp/train_resnet_fsdp_auto_wrap.py
     python3 examples/eager/train_decoder_only_eager.py