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

Remove unused param tpu_core_idx (#1948)

Browse Source
This commit is contained in:
Rohit Gupta
2020-05-25 16:04:53 -04:00
committed by GitHub
parent 34237cfcaf
commit d0ec11b9d6
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, tpu_core_idx, model):
def tpu_train(self, 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
@@ -900,7 +900,7 @@ class Trainer(
# train
if self.tpu_id is not None:
self.tpu_train(self.tpu_id, model)
self.tpu_train(model)
else:
xmp.spawn(self.tpu_train, args=(model,), nprocs=self.tpu_cores, start_method=start_method)