Fix failing test on main #3378

Job	Run time
run-tests (latest, test_prod)	4m 15s
run-tests (latest, test_core)	6m 23s
run-tests (latest, test_cli)	6m 15s
run-tests (latest, test_big_modeling)	6m 19s
run-tests (latest, test_deepspeed)	6m 42s
run-tests (latest, test_fsdp)	5m 33s
run-tests (latest, test_example_differences)	5m 23s
run-tests (latest, test_checkpoint_step)	9m 43s
run-tests (latest, test_checkpoint_epoch)	9m 7s
run-tests (latest, test_rest)	15m 57s
run-tests (minimum, test_prod)	4m 58s
run-tests (minimum, test_core)	5m 26s
run-tests (minimum, test_cli)	7m 9s
run-tests (minimum, test_big_modeling)	6m 19s
run-tests (minimum, test_deepspeed)	5m 32s
run-tests (minimum, test_fsdp)	5m 26s
run-tests (minimum, test_example_differences)	6m 29s
run-tests (minimum, test_checkpoint_step)	10m 6s
run-tests (minimum, test_checkpoint_epoch)	7m 31s
run-tests (minimum, test_rest)	16m 0s
	2h 30m 33s

Provide feedback