Tests

Fix `total_num_steps` #3556

Job	Run time
PyTest (3.10)	4m 51s
pre-commit	46s
PyTest (3.11)	4m 33s
docker-e2e-tests (118, 11.8.0, 3.10, 2.1.2, --extra-index-url https://download.pytorch.org/whl/cu...	30m 52s
docker-e2e-tests (121, 12.1.0, 3.10, 2.1.2, 1)	26m 8s
docker-e2e-tests (121, 12.1.0, 3.11, 2.2.1, 1)	28m 58s
	1h 36m 8s

Provide feedback