Repository: Lightning-AI/pytorch-lightning

Issue:
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
```

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
```



Discussion:

Resolving patch:
diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py
--- a/pytorch_lightning/trainer/distrib_parts.py
+++ b/pytorch_lightning/trainer/distrib_parts.py
@@ -489,7 +489,7 @@ def tpu_train(self, tpu_core_idx, model):
 
         # init 16 bit for TPU
         if self.precision == 16:
-            os.environ['XLA_USE_BF16'] = 1
+            os.environ['XLA_USE_BF16'] = str(1)
 
         m = f'INIT TPU local core: {self.tpu_local_core_rank}, ' \
             f'global rank: {self.tpu_global_core_rank}'


Test patch: