Fix test_prompt_filtering which got broken after a rebase.

igorts-git · igorts-git · commit 970808c9a7d8 · 2026-03-17T17:48:11.000-07:00
diff --git 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!

Original file line number	Diff line number	Diff line change
`@@ -331,6 +331,7 @@ def get_dataset_side_effect(model_tokenizer, config, data_dir, split, data_files`
`331`	`331`	`train_fraction=1.0,`
`332`	`332`	`num_epoch=1,`
`333`	`333`	`num_test_batches=1,`
	`334`	`+ test_batch_start_index=0,`
`334`	`335`	`)`
`335`	`336`
`336`	`337`	`# Patch everything!`