Merge pull request #210 from peregilk/patch-1

Update train.py
princeton-nlp · Nov 14, 2022 · 13361d0 · 13361d0
2 parents 511c99d + 5387856
commit 13361d0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/train.py b/train.py
@@ -203,6 +203,7 @@ def _setup_devices(self) -> "torch.device":
             device = torch.device("cpu")
             self._n_gpu = 0
         elif is_torch_tpu_available():
+            import torch_xla.core.xla_model as xm
             device = xm.xla_device()
             self._n_gpu = 0
         elif self.local_rank == -1: