# Buggy code
def tpu_train(self, tpu_core_idx, model):
    # put model on tpu
    self._device = (
        xm.xla_device(self.tpu_id) if self.tpu_id is not None else xm.xla_device()
    )
    model.to(self._device)

    # get the appropriate tpu ranks
    self.tpu_local_core_rank = xm.get_local_ordinal()
    self.tpu_global_core_rank = xm.get_ordinal()

    # avoid duplicating progress bar
    if self.tpu_global_core_rank != 0 and self.progress_bar_callback is not None:
        self.progress_bar_callback.disable()

    self.global_rank = self.tpu_local_core_rank
    rank_zero_only.rank = self.global_rank

    # CHOOSE OPTIMIZER
    # allow for lr schedulers as well
    self.optimizers, self.lr_schedulers, self.optimizer_frequencies = (
        self.init_optimizers(model)
    )

    # init 16 bit for TPU
    if self.precision == 16:
        os.environ["XLA_USE_BF16"] = str(1)

    log.info(
        f"INIT TPU local core: {self.tpu_local_core_rank},"
        f" global rank: {self.tpu_global_core_rank}"
    )

    # continue training routine
    self.run_pretrain_routine(model)

    # when training ends on these platforms dump weights to get out of the main process
    if self.on_colab_kaggle:
        self.save_spawn_weights(model)


# Traceback (AssertionError)
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
CUDA_VISIBLE_DEVICES: [0,1,2,3]
Traceback (most recent call last):
File "/opt/conda/lib/python3.6/runpy.py", line 193, in _run_module_as_main
"__main__", mod_spec)
File "/opt/conda/lib/python3.6/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/opt/conda/lib/python3.6/site-packages/pl_examples/basic_examples/gpu_template.py", line 80, in <module>
main(hyperparams)
File "/opt/conda/lib/python3.6/site-packages/pl_examples/basic_examples/gpu_template.py", line 41, in main
trainer.fit(model)
File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 860, in fit
self.barrier('fit_prepare_data')
File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 1261, in barrier
torch_distrib.barrier()
File "/opt/conda/lib/python3.6/site-packages/torch/distributed/distributed_c10d.py", line 1484, in barrier
_check_default_pg()
File "/opt/conda/lib/python3.6/site-packages/torch/distributed/distributed_c10d.py", line 187, in _check_default_pg
"Default process group is not initialized"
AssertionError: Default process group is not initialized

# Fixed code
def tpu_train(self, tpu_core_idx, model):
    # call setup after the ddp process has connected
    self.setup("fit")
    if self.is_function_implemented("setup", model):
        model.setup("fit")

    # put model on tpu
    self._device = (
        xm.xla_device(self.tpu_id) if self.tpu_id is not None else xm.xla_device()
    )
    model.to(self._device)

    # get the appropriate tpu ranks
    self.tpu_local_core_rank = xm.get_local_ordinal()
    self.tpu_global_core_rank = xm.get_ordinal()

    # avoid duplicating progress bar
    if self.tpu_global_core_rank != 0 and self.progress_bar_callback is not None:
        self.progress_bar_callback.disable()

    self.global_rank = self.tpu_local_core_rank
    rank_zero_only.rank = self.global_rank

    # CHOOSE OPTIMIZER
    # allow for lr schedulers as well
    self.optimizers, self.lr_schedulers, self.optimizer_frequencies = (
        self.init_optimizers(model)
    )

    # init 16 bit for TPU
    if self.precision == 16:
        os.environ["XLA_USE_BF16"] = str(1)

    log.info(
        f"INIT TPU local core: {self.tpu_local_core_rank},"
        f" global rank: {self.tpu_global_core_rank}"
    )

    # continue training routine
    self.run_pretrain_routine(model)

    # when training ends on these platforms dump weights to get out of the main process
    if self.on_colab_kaggle:
        self.save_spawn_weights(model)