Run pytorch code on Kaggle TPUs

Hi All,

I have some code. Which I borrowed from here. Just wondering how to edit it so that I would be able to utilise the TPUs in Kaggle notebooks.

Would anyone be able to help me in this regards.

Thanks & Best Regards