TPUs: crash using torch-xla nightly
🐛 Bug
If I try to use torch-xla nightly with PyTorch Lightning, I see this crash:
(torch-xla-nightly) zcain@zcain-pl-verify:~/pytorch-lightning/pl_examples/domain_templates$ python computer_vision_fine_tuning.py 
Traceback (most recent call last):
  File "computer_vision_fine_tuning.py", line 55, in <module>
    import pytorch_lightning as pl
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/pytorch_lightning/__init__.py", line 65, in <module>
    from pytorch_lightning.trainer import Trainer
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/pytorch_lightning/trainer/__init__.py", line 18, in <module>
    from pytorch_lightning.trainer.trainer import Trainer
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 30, in <module>
    from pytorch_lightning.loggers import LightningLoggerBase
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/pytorch_lightning/loggers/__init__.py", line 18, in <module>
    from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/pytorch_lightning/loggers/tensorboard.py", line 24, in <module>
    from torch.utils.tensorboard import SummaryWriter
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/utils/tensorboard/__init__.py", line 8, in <module>
    from .writer import FileWriter, SummaryWriter  # noqa F401
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/utils/tensorboard/writer.py", line 9, in <module>
    from tensorboard.compat.proto.event_pb2 import SessionLog
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/tensorboard/compat/proto/event_pb2.py", line 17, in <module>
    from tensorboard.compat.proto import summary_pb2 as tensorboard_dot_compat_dot_proto_dot_summary__pb2
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/tensorboard/compat/proto/summary_pb2.py", line 17, in <module>
    from tensorboard.compat.proto import tensor_pb2 as tensorboard_dot_compat_dot_proto_dot_tensor__pb2
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/tensorboard/compat/proto/tensor_pb2.py", line 16, in <module>
    from tensorboard.compat.proto import resource_handle_pb2 as tensorboard_dot_compat_dot_proto_dot_resource__handle__pb2
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/tensorboard/compat/proto/resource_handle_pb2.py", line 150, in <module>
    '__module__' : 'tensorboard.compat.proto.resource_handle_pb2'
SystemError: google/protobuf/pyext/descriptor.cc:354: bad argument to internal function

Repro
create VM:
gcloud compute instances create zcain-vm
--zone=us-central1-a
--machine-type=e2-highmem-16
--image-family=torch-xla
--image-project=ml-images
--boot-disk-size=300GB
--scopes=https://www.googleapis.com/auth/cloud-platform
create TPU:
gcloud compute tpus create zcain-tpu
--zone=us-central1-a
--network=default
--version=pytorch-nightly
--accelerator-type=v3-8
SSH into VM:
(VM) git clone https://github.com/PyTorchLightning/pytorch-lightning.git
(VM) conda activate torch-xla-nightly
(VM) export TPU_IP_ADDRESS=<TPU's IP address>
(VM) export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
(VM) cd pytorch-lightning
(VM) pip install .
(VM) pip install -r requirements/test.txt
(VM) cd pl_examples/domain_templates/
(VM) vim computer_vision_fine_tuning.py
(VM) <replace GPUs arg with tpu_cores=8,>
(VM) python computer_vision_fine_tuning.py

Note: I get the same crash when using pip install pytorch-lightning instead of installing from source
If I use torch-xla-1.7, training works fine.
Environment
* CUDA:
        - GPU:
        - available:         False
        - version:           None
* Packages:
        - numpy:             1.19.2
        - pyTorch_debug:     False
        - pyTorch_version:   1.9.0a0+8c185e6
        - pytorch-lightning: 1.1.8
        - tqdm:              4.56.0
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                - 
        - processor:         
        - python:            3.6.10
        - version:           #1 SMP Debian 4.9.246-2 (2020-12-17)

Additional context
The crash seems related to Tensorboard. I wonder if the issue is related to Tensorboard versioning, however pip install pytorch-lightning doesn't seem to install any new tensorboard version compared to what is already installed in this public torch-xla conda env:
Installing collected packages: PyYAML, fsspec, pytorch-lightning, yarl
  Attempting uninstall: PyYAML
    Found existing installation: PyYAML 5.4.1
    Uninstalling PyYAML-5.4.1:
      Successfully uninstalled PyYAML-5.4.1
  Attempting uninstall: fsspec
    Found existing installation: fsspec 0.7.4
    Uninstalling fsspec-0.7.4:
      Successfully uninstalled fsspec-0.7.4
  Attempting uninstall: yarl
    Found existing installation: yarl 1.6.3
    Uninstalling yarl-1.6.3:
      Successfully uninstalled yarl-1.6.3
Successfully installed PyYAML-5.3.1 fsspec-0.8.5 pytorch-lightning-1.1.8 yarl-1.5.1

Let me know if you can think of anything else I should try