Would be nice if this repo supported training on TPUs.
Would be nice if this repo supported training on TPUs.