Fix #997 (#1018)

2020-03-02 20:51:05 -06:00 · 2020-03-02 20:51:05 -06:00 · 29cbc9e723
parent 6dae5698ef
commit 29cbc9e723
1 changed files with 1 additions and 1 deletions
--- a/pytorch_lightning/trainer/distrib_parts.py
+++ b/pytorch_lightning/trainer/distrib_parts.py
@ -489,7 +489,7 @@ class TrainerDPMixin(ABC):

        # 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}'