Repository: Lightning-AI/pytorch-lightning

Issue:
RuntimeError: Unimplemented backend XLA on TPU
## 🐛 Bug

`RuntimeError: Unimplemented backend XLA` raised for `self.batch_loss_value.append(loss)` line in `trainer/training_loop.py` file when running MNIST on TPU. I think it was introduced in 31b7148.

### To Reproduce

Steps to reproduce the behavior:

1. Go to [MNIST on TPUs](https://colab.research.google.com/drive/1-_LKx4HwAxl5M6xPJmqAAu444LTDQoa3)
2. Run all
3. Scroll down to trainer
4. See error

```
Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py", line 119, in _start_fn
    fn(gindex, *args)
  File "/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py", line 119, in _start_fn
    fn(gindex, *args)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/distrib_parts.py", line 499, in tpu_train
    self.run_pretrain_routine(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/distrib_parts.py", line 499, in tpu_train
    self.run_pretrain_routine(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/distrib_parts.py", line 499, in tpu_train
    self.run_pretrain_routine(model)
  File "/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py", line 119, in _start_fn
    fn(gindex, *args)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/distrib_parts.py", line 499, in tpu_train
    self.run_pretrain_routine(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/distrib_parts.py", line 499, in tpu_train
    self.run_pretrain_routine(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/distrib_parts.py", line 499, in tpu_train
    self.run_pretrain_routine(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 920, in run_pretrain_routine
    self.train()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 920, in run_pretrain_routine
    self.train()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 920, in run_pretrain_routine
    self.train()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 920, in run_pretrain_routine
    self.train()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 920, in run_pretrain_routine
    self.train()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 920, in run_pretrain_routine
    self.train()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 356, in train
    self.run_training_epoch()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 356, in train
    self.run_training_epoch()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/distrib_parts.py", line 499, in tpu_train
    self.run_pretrain_routine(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 356, in train
    self.run_training_epoch()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 356, in train
    self.run_training_epoch()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 356, in train
    self.run_training_epoch()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 356, in train
    self.run_training_epoch()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 425, in run_training_epoch
    output = self.run_training_batch(batch, batch_idx)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 425, in run_training_epoch
    output = self.run_training_batch(batch, batch_idx)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 920, in run_pretrain_routine
    self.train()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 425, in run_training_epoch
    output = self.run_training_batch(batch, batch_idx)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 425, in run_training_epoch
    output = self.run_training_batch(batch, batch_idx)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 425, in run_training_epoch
    output = self.run_training_batch(batch, batch_idx)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 425, in run_training_epoch
    output = self.run_training_batch(batch, batch_idx)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 582, in run_training_batch
    self.batch_loss_value.append(loss)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 582, in run_training_batch
    self.batch_loss_value.append(loss)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/supporting_classes.py", line 23, in append
    if self.memory.type() != x.type():
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 356, in train
    self.run_training_epoch()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 582, in run_training_batch
    self.batch_loss_value.append(loss)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 582, in run_training_batch
    self.batch_loss_value.append(loss)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 582, in run_training_batch
    self.batch_loss_value.append(loss)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/supporting_classes.py", line 23, in append
    if self.memory.type() != x.type():
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 582, in run_training_batch
    self.batch_loss_value.append(loss)
Exception in device=TPU:5: Unimplemented backend XLA
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/supporting_classes.py", line 23, in append
    if self.memory.type() != x.type():
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 425, in run_training_epoch
    output = self.run_training_batch(batch, batch_idx)
RuntimeError: Unimplemented backend XLA
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/supporting_classes.py", line 23, in append
    if self.memory.type() != x.type():
RuntimeError: Unimplemented backend XLA
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/supporting_classes.py", line 23, in append
    if self.memory.type() != x.type():
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/supporting_classes.py", line 23, in append
    if self.memory.type() != x.type():
RuntimeError: Unimplemented backend XLA
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 582, in run_training_batch
    self.batch_loss_value.append(loss)
RuntimeError: Unimplemented backend XLA
Traceback (most recent call last):
RuntimeError: Unimplemented backend XLA
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/supporting_classes.py", line 23, in append
    if self.memory.type() != x.type():
RuntimeError: Unimplemented backend XLA
  File "/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py", line 119, in _start_fn
    fn(gindex, *args)
RuntimeError: Unimplemented backend XLA
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/distrib_parts.py", line 499, in tpu_train
    self.run_pretrain_routine(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 920, in run_pretrain_routine
    self.train()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 356, in train
    self.run_training_epoch()
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 425, in run_training_epoch
    output = self.run_training_batch(batch, batch_idx)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 582, in run_training_batch
    self.batch_loss_value.append(loss)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/supporting_classes.py", line 23, in append
    if self.memory.type() != x.type():
RuntimeError: Unimplemented backend XLA
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-2-12f6e300d51d> in <module>()
      6 # most basic trainer, uses good defaults
      7 trainer = Trainer(num_tpu_cores=8)
----> 8 trainer.fit(model)

3 frames
/usr/local/lib/python3.6/dist-packages/torch/multiprocessing/spawn.py in join(self, timeout)
    111                 raise Exception(
    112                     "process %d terminated with exit code %d" %
--> 113                     (error_index, exitcode)
    114                 )
    115 

Exception: process 4 terminated with exit code 17
```

### Environment

 - PyTorch Version (e.g., 1.0): 1.6
 - OS (e.g., Linux): Linux
 - How you installed PyTorch (`conda`, `pip`, source): pip
 - Build command you used (if compiling from source): -
 - Python version: 3.6
 - CUDA/cuDNN version: -
 - GPU models and configuration: -
 - Any other relevant information: TPU backend



Discussion:
@williamFalcon 
The issue is caused by [this line](https://github.com/PyTorchLightning/pytorch-lightning/blob/f1e11d8b3874067016693c50ae253ec79eecda09/pytorch_lightning/trainer/supporters.py#L40):
```
        if self.memory.type() != x.type():
            self.memory.type_as(x)
```
For TPU `x` is a XLA tensor and `x.type()` results in `Unimplemented backend XLA` (for GPU type is `torch.cuda.FloatTensor`).

Something like `x = torch.Tensor([x])` before condition checking fixes the problem. Or we can just send one of the tensors to common device.

Notebook for debugging on [Google Colab](https://colab.research.google.com/drive/16Ug8IYPkqCu_NhK1FV7W-vszDsdgXfPD).

Resolving patch:
diff --git a/pytorch_lightning/trainer/supporters.py b/pytorch_lightning/trainer/supporters.py
--- a/pytorch_lightning/trainer/supporters.py
+++ b/pytorch_lightning/trainer/supporters.py
@@ -36,9 +36,9 @@ def last(self):
             return self.memory[self.last_idx]
 
     def append(self, x):
-        # map proper type for memory if they don't match
-        if self.memory.type() != x.type():
-            self.memory.type_as(x)
+        # ensure same device and type
+        if self.memory.device != x.device or self.memory.type() != x.type():
+            x = x.to(self.memory)
 
         # store without grads
         with torch.no_grad():


Test patch: