How to run GPT Neo on TPU using transformer?

I’ve run GPT Neo on TPU…
It’s stucked in model.generate()

What’s the proper way to run it on TPU?