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

This reverts commit d0ec11b9d6.

pytorch_lightning/trainer/distrib_parts.py

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

@@ -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)