Loss does not decrease on single TPU

https://github.com/pytorch/xla/issues/2735#issue-786787271

🐛 Bug

With pytorch-lightning, the trainer hangs up at the end of epoch. By keyboard interruption, it seems that it hangs up at fd_event_list = self._poll.poll(timeout)

---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
<ipython-input-4-fd5234884739> in <module>()
      4 autoencoder_1 = LitAutoEncoder()
      5 trainer_1 = pl.Trainer(tpu_cores=8, max_epochs=1, progress_bar_refresh_rate=20)
----> 6 trainer_1.fit(autoencoder_1, DataLoader(dataset_mnist_train))

6 frames
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
    471         self.call_hook('on_fit_start')
    472 
--> 473         results = self.accelerator_backend.train()
    474         self.accelerator_backend.teardown()
    475 

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/accelerators/tpu_accelerator.py in train(self)
    111                 args=(model, self.trainer, self.mp_queue),
    112                 nprocs=self.trainer.tpu_cores,
--> 113                 start_method=self.start_method
    114             )
    115 

/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py in spawn(fn, args, nprocs, join, daemon, start_method)
    393         join=join,
    394         daemon=daemon,
--> 395         start_method=start_method)
    396 
    397 

/usr/local/lib/python3.6/dist-packages/torch/multiprocessing/spawn.py in start_processes(fn, args, nprocs, join, daemon, start_method)
    155 
    156     # Loop on join until it returns True or raises an exception.
--> 157     while not context.join():
    158         pass
    159 

/usr/local/lib/python3.6/dist-packages/torch/multiprocessing/spawn.py in join(self, timeout)
     75         ready = multiprocessing.connection.wait(
     76             self.sentinels.keys(),
---> 77             timeout=timeout,
     78         )
     79 

/usr/lib/python3.6/multiprocessing/connection.py in wait(object_list, timeout)
    909 
    910             while True:
--> 911                 ready = selector.select(timeout)
    912                 if ready:
    913                     return [key.fileobj for (key, events) in ready]

/usr/lib/python3.6/selectors.py in select(self, timeout)
    374             ready = []
    375             try:
--> 376                 fd_event_list = self._poll.poll(timeout)
    377             except InterruptedError:
    378                 return ready

KeyboardInterrupt: 

To Reproduce

Steps to reproduce the behavior:

1. Run this Colab (or another Colab for issue reporting) with TPU.
2. Wait until the first (and the last) epoch finishes. (all steps are completed but it keeps running)
3. Interrupt kernel

Expected behavior

The training should be completed immediately after all steps are completed.

Environment

• Reproducible on XLA backend [CPU/TPU]: TPU
• torch_xla version: 1.7 and nightly

cc @carmocca @JackCaoG @Liyang90 @gkroiz @kaushikb11 @rohitgr7 @akihironitta