# Buggy code
def save_checkpoint(self, filepath, weights_only: bool = False):
    """Save model/training states as a checkpoint file through state-dump and file-write.

    Args:
        filepath: write-target file's path
        weights_only: saving model weights only
    """
    # dump states as a checkpoint dictionary object
    _checkpoint = self.lightning_module.trainer.checkpoint_connector.dump_checkpoint(
        weights_only
    )
    # Todo: TypeError: 'mappingproxy' object does not support item assignment
    xm.save({k: v for k, v in _checkpoint.items() if k != "callbacks"}, filepath)


# Traceback (RuntimeError)
RuntimeError                              Traceback (most recent call last)
<ipython-input-17-587e2a9e3858> in <module>
----> 1 trainer.test(datamodule=dm)

/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in test(self, model, test_dataloaders, ckpt_path, verbose, datamodule)
922             results = self.__test_given_model(model, test_dataloaders)
923         else:
--> 924             results = self.__test_using_best_weights(ckpt_path, test_dataloaders)
925
926         self.teardown('test')

/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in __test_using_best_weights(self, ckpt_path, test_dataloaders)
950                 return {}
951             if not self._device_type == DeviceType.TPU:
--> 952                 self.accelerator.barrier()
953
954             ckpt = pl_load(ckpt_path, map_location=lambda storage, loc: storage)

/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/accelerator.py in barrier(self, name)
375
376     def barrier(self, name: Optional[str] = None) -> None:
--> 377         self.training_type_plugin.barrier(name=name)
378
379     def broadcast(self, obj: object, src: int = 0) -> object:

/opt/conda/lib/python3.7/site-packages/pytorch_lightning/plugins/training_type/tpu_spawn.py in barrier(self, name)
113
114     def barrier(self, name: Optional[str] = None) -> None:
--> 115         rendezvous(f"pl.Trainer.{name}")
116
117     def transfer_distrib_spawn_state_on_fit_end(self, results):

/opt/conda/lib/python3.7/site-packages/torch_xla/core/xla_model.py in rendezvous(tag, payload, replicas)
859     ordinal `i` at position `i` in the returned tuple.
860   """
--> 861   return torch_xla._XLAC._xla_rendezvous(get_ordinal(), tag, payload, replicas)
862
863

RuntimeError: tensorflow/compiler/xla/xla_client/mesh_service.cc:316 : Check failed: impl_->channel->WaitForConnected( std::chrono::system_clock::now() + std::chrono::seconds(connect_wait_seconds))

# Fixed code
def save_checkpoint(self, filepath, weights_only: bool = False):
    """Save model/training states as a checkpoint file through state-dump and file-write.

    Args:
        filepath: write-target file's path
        weights_only: saving model weights only
    """
    # dump states as a checkpoint dictionary object
    _checkpoint = self.lightning_module.trainer.checkpoint_connector.dump_checkpoint(
        weights_only
    )
    # Todo: TypeError: 'mappingproxy' object does not support item assignment
    self.save({k: v for k, v in _checkpoint.items() if k != "callbacks"}, filepath)