Add fsdpv2 test to tpu ci

pytorch · May 10, 2024 · 76c7435 · 76c7435
1 parent aeb59b5
commit 76c7435
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/test/tpu/run_tests.sh b/test/tpu/run_tests.sh
@@ -11,6 +11,7 @@ python3 test/spmd/test_xla_distributed_checkpoint.py
 python3 test/spmd/test_train_spmd_linear_model.py
 python3 test/spmd/test_xla_spmd_python_api_interaction.py
 python3 test/spmd/test_xla_auto_sharding.py
+python3 test/spmd/test_fsdp_v2.py
 XLA_EXPERIMENTAL=nonzero:masked_select:nms python3 test/ds/test_dynamic_shape_models.py -v
 XLA_EXPERIMENTAL=nonzero:masked_select:nms python3 test/ds/test_dynamic_shapes.py -v
 python3 test/test_autocast.py