add pallas_spmd.py to tpu ci

pytorch · May 10, 2024 · c28640c · c28640c
1 parent 887d344
commit c28640c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/test/tpu/run_tests.sh b/test/tpu/run_tests.sh
@@ -21,6 +21,7 @@ python3 test/pjrt/test_dtypes.py
 python3 test/pjrt/test_dynamic_plugin_tpu.py
 python3 test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py
 python3 test/test_pallas.py
+python3 test/test_pallas_spmd.py
 python3 test/test_input_output_aliases.py
 python3 test/torch_distributed/test_torch_distributed_all_gather_xla_backend.py
 python3 test/torch_distributed/test_torch_distributed_all_reduce_xla_backend.py