Repository: pytorch
Issue: Installed successfully the torch-1.9 and torch_xla-1.9. RuntimeError requested XLA_GPU:0，but available devices are [CPU:0, XLA_CPU:0 ]
Labels: module: binaries,triaged,module: xla

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:

1. running in the docker
2.pip install torch==1.9
3.pip install https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp36-cp36m-linux_x86_64.whl
4. Run the XLA test case: python /xla/test/test_train_mp_mnist.py

## Expected behavior
Expected  no stack exceptions. but i got the error: 
Traceback (most recent call last):
  File "test_train_mp_mnist.py", line 200, in <module>
    xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=FLAGS.num_cores)
  File "/root/miniconda3/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 386, in spawn
    _start_fn(0, pf_cfg, fn, args)
  File "/root/miniconda3/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 323, in _start_fn
    fn(gindex, *args)
  File "test_train_mp_mnist.py", line 190, in _mp_fn
    accuracy = train_mnist(flags)
  File "test_train_mp_mnist.py", line 126, in train_mnist
    model = MNIST().to(device)
  File "/root/miniconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 852, in to
    return self._apply(convert)
  File "/root/miniconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 530, in _apply
    module._apply(fn)
  File "/root/miniconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 552, in _apply
    param_applied = fn(param)
  File "/root/miniconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 850, in convert
    return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
RuntimeError: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:421 : Check failed: session->session()->Run(session_work->feed_inputs, session_work->outputs_handles, &outputs) ==us::OK() (Invalid argument: Cannot assign a device for operation XRTAllocateFromTensor: {{node XRTAllocateFromTensor}} was explicitly assigned to /job:localservice/replica:0/task:0/device:Xble devices are [ /job:localservice/replica:0/task:0/device:CPU:0, /job:localservice/replica:0/task:0/device:XLA_CPU:0 ]. Make sure the device specification refers to a valid device.
         [[XRTAllocateFromTensor]] vs. OK)
*** Begin stack trace ***
        tensorflow::CurrentStackTrace()
        xla::util::MultiWait::Complete(std::function<void ()> const&)

        clone
*** End stack trace ***

## Environment
Collecting environment information...
PyTorch version: 1.9.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.21.1
Libc version: glibc-2.9

Python version: 3.6.5 |Anaconda, Inc.| (default, Mar 29 2018, 18:21:58)  [GCC 7.2.0] (64-bit runtime)
Python platform: Linux-4.4.0-21-generic-x86_64-with-debian-buster-sid
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration:
GPU 0: NVIDIA Tesla V100-SXM2-32GB
GPU 1: NVIDIA Tesla V100-SXM2-32GB
GPU 2: NVIDIA Tesla V100-SXM2-32GB
GPU 3: NVIDIA Tesla V100-SXM2-32GB

Nvidia driver version: 465.19.01
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.15.4
[pip3] torch==1.9.0
[pip3] torch-xla==1.9
[pip3] torchvision==0.10.0
[conda] mkl-include               2021.3.0                  <pip>
[conda] numpy                     1.15.4                    <pip>
[conda] torch                     1.9.0                     <pip>
[conda] torch-xla                 1.9                       <pip>
[conda] torchvision               0.10.0                    <pip>


cc @ezyang @seemethere @malfet @walterddr @bdhirsh @gmagogsfm