Precision=16 with TPUs bug
🐛 Bug
Setting precision=16 when training with a TPU throws an error
To Reproduce
see colab: https://colab.research.google.com/drive/1s-ZDIqzgKQ1Byf-Lw58RZ8LGgmdB6qjB
Relavent stack trace:
Exception in device=TPU:0: str expected, not int
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/pytorch_lightning/trainer/distrib_parts.py", line 492, in tpu_train
    os.environ['XLA_USE_BF16'] = 1
  File "/usr/lib/python3.6/os.py", line 674, in __setitem__
    value = self.encodevalue(value)
  File "/usr/lib/python3.6/os.py", line 744, in encode
    raise TypeError("str expected, not %s" % type(value).__name__)
TypeError: str expected, not int

To fix this all that should be needed should be casting 1 to a string
    os.environ['XLA_USE_BF16'] = str(1)

Environment
Collecting environment information...
PyTorch version: 1.4.0
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04.3 LTS
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
CMake version: version 3.12.0

Python version: 3.6
Is CUDA available: No
CUDA runtime version: 10.1.243
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.17.5
[pip3] torch==1.4.0
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.3.1
[pip3] torchvision==0.5.0
[conda] Could not collect