specifying the tpu_core speed-up TPU training

[rohitgr7]

🐛 Bug

1. I am getting a huge time difference between training a model on a specific tpu core tpu_cores=[1] and training a model on just 1 tpu core tpu_cores=1. What's the reason for that? Aren't both the conditions the same with just the difference that I am assigning a specific tpu_core in the first case and assigning the number of tpu_cores I want to use in the second case. Also in the second case, I am getting an error. When training with tpu_cores=[1] epoch time is 17 seconds with tpu_cores=1 epoch time is just 5 seconds.
2. Running on colab gives me an error but no error on Kaggle kernels. But the time difference issue is the same on both the platforms.

To Reproduce

Code sample

Colab Notebook

Expected behavior

As far as I know in both cases, the training time should be the same regardless of training on a single core or training on a specific core.

Environment

• PyTorch Version (e.g., 1.0): 1.5.0
• OS (e.g., Linux): Linux
• How you installed PyTorch (conda, pip, source): pip
• Build command you used (if compiling from source):
• Python version: 3.7
• CUDA/cuDNN version: 10.1
• GPU models and configuration: Tesla P100-PCIE-16GB
• Any other relevant information:

Additional context

[rohitgr7]

@williamFalcon @dlibenzi

[dlibenzi]

Because of a bug?

https://github.com/PyTorchLightning/pytorch-lightning/blob/fdbbe968256f6c68a5dbb840a2004b77a618ef61/pytorch_lightning/trainer/trainer.py#L360

https://github.com/PyTorchLightning/pytorch-lightning/blob/fdbbe968256f6c68a5dbb840a2004b77a618ef61/pytorch_lightning/trainer/training_loop.py#L415

[lezwon]

@dlibenzi I recall when training on a single core and using ParallelLoader, I used to receive an error. Hence the self.tpu_id is None condition. However, I did a recheck and it seems to be working fine with ParallelLoader now. Made a PR for the same. :)

[Borda]

not sure why this was reopened...