diff --git a/tests/post_training/unit/train_rl_test.py b/tests/post_training/unit/train_rl_test.py
index e5a4aa8ec2..4bb9831f60 100644
--- a/tests/post_training/unit/train_rl_test.py
+++ b/tests/post_training/unit/train_rl_test.py
@@ -331,6 +331,7 @@ def get_dataset_side_effect(model_tokenizer, config, data_dir, split, data_files
         train_fraction=1.0,
         num_epoch=1,
         num_test_batches=1,
+        test_batch_start_index=0,
     )
 
     # Patch everything!