wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

Revert "Remove unused param tpu_core_idx (#1948)" (#1963)

Browse Source 
This reverts commit d0ec11b9d6.
This commit is contained in:
Jirka Borovec
2020-05-26 19:02:51 -04:00
committed by GitHub
parent 460ab5485e
commit ca815698f5
2 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/trainer/distrib_parts.py
+1 -1
View File
@@ -501,7 +501,7 @@ class TrainerDPMixin(ABC):
self.run_pretrain_routine(model)
def tpu_train(self, model):
def tpu_train(self, tpu_core_idx, model):
# put model on tpu
self._device = xm.xla_device(self.tpu_id) if self.tpu_id is not None else xm.xla_device()
model.to(self._device)
pytorch_lightning/trainer/trainer.py
+1 -1
View File
@@ -904,7 +904,7 @@ class Trainer(
# train
if self.tpu_id is not None:
self.tpu_train(model)
self.tpu_train(self.tpu_id, model)
else:
xmp.spawn(self.tpu_train, args=(model,), nprocs=self.tpu_cores, start_method=start_method)