Revert "Remove unused param tpu_core_idx (#1948)"

This reverts commit d0ec11b.

Author: Borda

pytorch_lightning/trainer/distrib_parts.py

@@ -501,7 +501,7 @@ def single_gpu_train(self, model):
 
         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

@@ -900,7 +900,7 @@ def fit(
 
             # 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)