Fix UT axlearn.common.t5_test

Author: camiloCienet

axlearn/common/t5_test.py

@@ -342,7 +342,7 @@ def test_against_t5_encoder_decoder_model(self, arch):
     def test_pjit(self):
         # A simple test to ensure a train step does not leak tracers.
         mesh_shape = (1, 1)
-        mesh_axes = ("data", "model")
+        mesh_axes = ("fsdp", "model")
         devices = mesh_utils.create_device_mesh(mesh_shape)
         with jax.checking_leaks(), jax.sharding.Mesh(devices, mesh_axes):
             vocab_size = 6