fix: support multi-gpu training #1729

Job	Run time
Documentation	3m 35s
Linting	0s
Pytest (3.7)	12m 58s
Pytest (3.8)	12m 8s
Simple installation (3.7)	1m 28s
Pytest (3.9)	20m 49s
Simple installation (3.8)	1m 42s
Simple installation (3.9)	1m 31s
Simple installation (3.10)	2m 3s
Simple installation (3.11)	1m 50s
Simple installation (3.12)	1m 50s
Coverage / Check Pull Request Coverage	0s
Coverage / Update Base Coverage	13s
	1h 0m 7s

Provide feedback