chore: refactor training loop (#4435) #3492

Job	Run time
Build source distribution	11s
determine-arm64-runner	1s
Build wheels for cp311-manylinux_x86_64	6m 33s
Build wheels for cp311-manylinux_x86_64	6m 22s
Build wheels for cp311-macosx_x86_64	6m 29s
Build wheels for cp311-macosx_arm64	2m 41s
Build wheels for cp311-win_amd64	5m 36s
Build wheels for cp310-manylinux_aarch64	30m 56s
build_docker (12)	2m 24s
build_docker (_cu11, 11)	2m 53s
build_pypi_index	9s
upload_pypi	0s
deploy_pypi_index	0s
Pass testing build wheels	2s
	1h 4m 17s

Provide feedback