Make torch xla available on GPU #3587

Job	Run time
run-tests (latest, test_prod)	3m 39s
run-tests (latest, test_core)	4m 33s
run-tests (latest, test_cli)	5m 46s
run-tests (latest, test_big_modeling)	4m 45s
run-tests (latest, test_deepspeed)	5m 6s
run-tests (latest, test_fsdp)	4m 56s
run-tests (latest, test_example_differences)	4m 27s
run-tests (latest, test_checkpoint_step)	6m 30s
run-tests (latest, test_checkpoint_epoch)	16m 28s
run-tests (latest, test_rest)	10m 26s
run-tests (minimum, test_prod)	3m 42s
run-tests (minimum, test_core)	4m 44s
run-tests (minimum, test_cli)	5m 50s
run-tests (minimum, test_big_modeling)	4m 49s
run-tests (minimum, test_deepspeed)	4m 32s
run-tests (minimum, test_fsdp)	4m 45s
run-tests (minimum, test_example_differences)	4m 31s
run-tests (minimum, test_checkpoint_step)	6m 6s
run-tests (minimum, test_checkpoint_epoch)	14m 5s
run-tests (minimum, test_rest)	10m 17s
	2h 9m 57s

Provide feedback