CRITICAL: fix failing ci #3384

Job	Run time
run-tests (latest, test_prod)	4m 0s
run-tests (latest, test_core)	5m 58s
run-tests (latest, test_cli)	8m 6s
run-tests (latest, test_big_modeling)	5m 52s
run-tests (latest, test_deepspeed)	6m 39s
run-tests (latest, test_fsdp)	7m 39s
run-tests (latest, test_example_differences)	6m 29s
run-tests (latest, test_checkpoint_step)	9m 12s
run-tests (latest, test_checkpoint_epoch)	8m 17s
run-tests (latest, test_rest)	10m 27s
run-tests (minimum, test_prod)	4m 0s
run-tests (minimum, test_core)	5m 30s
run-tests (minimum, test_cli)	7m 59s
run-tests (minimum, test_big_modeling)	5m 49s
run-tests (minimum, test_deepspeed)	5m 10s
run-tests (minimum, test_fsdp)	5m 46s
run-tests (minimum, test_example_differences)	5m 10s
run-tests (minimum, test_checkpoint_step)	8m 19s
run-tests (minimum, test_checkpoint_epoch)	10m 33s
run-tests (minimum, test_rest)	16m 1s
	2h 26m 56s

Provide feedback