I want to signal-boost this. TPU support on PyTorch is partial. You can run modeling computation on the TPU with PyTorch, but not the data-loading. And without the TPU's data-loading, you're significantly, significantly bottle-necked, to the point where you are often better off using GPUs. The reason why TensorFlow and TPUs synergize so well is that the TPUs themselves can consume data for training, allowing for massive scalability.

I have great respect for the PyTorch-TPU team, but I would recommend not heavily advertising PyTorch-TPU support until this major feature disparity is made up.