Running PyTorch on TPU: a bag of tricks

At the time of writing these lines running PyTorch code on TPUs is not a well-trodden path. Naturally, TPUs have been optimized for and mainly used with TensorFlow. But Kaggle and Google distribute free TPU time on some of its competitions, and one doesn't simply change his favorite framework, so this is a memo on my (mostly successful) experience of training PyTorch models with TPU on GCP. PyTorch/XLA is the project that allows doing it. It is still in active development, issues get fixed.