Better error when device mismatches when calling gather() on CUDA #3566

Job	Run time
run-tests (latest, test_prod)	2m 58s
run-tests (latest, test_core)	4m 16s
run-tests (latest, test_cli)	5m 22s
run-tests (latest, test_big_modeling)	4m 7s
run-tests (latest, test_deepspeed)	4m 10s
run-tests (latest, test_fsdp)	4m 13s
run-tests (latest, test_example_differences)	3m 58s
run-tests (latest, test_checkpoint_step)	5m 41s
run-tests (latest, test_checkpoint_epoch)	5m 31s
run-tests (latest, test_rest)	9m 28s
run-tests (minimum, test_prod)	2m 59s
run-tests (minimum, test_core)	4m 34s
run-tests (minimum, test_cli)	5m 17s
run-tests (minimum, test_big_modeling)	4m 11s
run-tests (minimum, test_deepspeed)	4m 15s
run-tests (minimum, test_fsdp)	4m 3s
run-tests (minimum, test_example_differences)	4m 7s
run-tests (minimum, test_checkpoint_step)	5m 50s
run-tests (minimum, test_checkpoint_epoch)	5m 13s
run-tests (minimum, test_rest)	9m 46s
	1h 39m 59s

Provide feedback