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:

Go to MNIST on TPUs
Run all
Scroll down to trainer
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