TPU in Colab?

[AffeZwei]

Hi everyone,
have any of you been able to adapt the code so that the Colab version can run with TPU? At the moment, the train.py code only enables GPU, and a quick Google search shows that running a TPU with pytorch would require much more coding skill that I have at the moment - and all my attempts have failed. It's not just a question of changing the runtime type in the menu. I have ColabPro+, so I should be able to run with TPU, theoretically. Has anyone been able to implement this, and would be willing to share their code?