Comment by jspisak

We have been working with G on TPU support for PyTorch. Have you tried it out? https://github.com/pytorch/xla