Unless you physically work next to the TPU hardware team, the torch support for TPUs is pretty brittle.