training: Linting, fixes and docstrings #314

Job	Run time
pip (ubuntu-latest, cpu)	28s
pip (windows-latest, cpu)	1m 4s
conda (ubuntu-latest, nvidia)	8m 36s
pip (windows-latest, directml)	1m 0s
conda (ubuntu-latest, cpu)	8m 44s
conda (macos-latest, nvidia)	12m 32s
conda (macos-latest, cpu)	13m 45s
conda (windows-latest, nvidia)	13m 32s
conda (windows-latest, cpu)	15m 19s
conda (ubuntu-latest, rocm)	8m 49s
conda (windows-latest, directml)	14m 43s
	1h 38m 32s

Provide feedback