fixes in ltx_transformer_step_test.py

Author: prishajain1

src/maxdiffusion/tests/ltx_transformer_step_test.py

@@ -153,7 +153,7 @@ def test_one_step_transformer(self):
     state_shardings["transformer"] = transformer_state_shardings
     states["transformer"] = transformer_state
     example_inputs = {}
-    batch_size, num_tokens = 4, 256
+    batch_size, num_tokens = 8, 256
     input_shapes = {
         "latents": (batch_size, num_tokens, in_channels),
         "fractional_coords": (batch_size, 3, num_tokens),
@@ -194,6 +194,10 @@ def test_one_step_transformer(self):
 
     noise_pred = p_run_inference(states).block_until_ready()
     noise_pred = torch.from_numpy(np.array(noise_pred))
+    
+    # Using batch_size=8 to satisfy fsdp=8 sharding constraint, but reference is batch_size=4
+    if batch_size != noise_pred_pt.shape[0]:
+      noise_pred = noise_pred[:noise_pred_pt.shape[0]]
 
     torch.testing.assert_close(noise_pred_pt, noise_pred, atol=0.025, rtol=20)