Issue with using TPU of Unet model

[yashchks87]

Hello,

I tried using your model it works perfectly fine with GPU but when I try to use it with TPU in Colab and also GCP VM with TPU same error comes. Error name is Socket Closed. Let me know if you need any more explanation.

[bomerzz]

Try lowering your input size or batch size. Socket closed normally means you are overloading the TPU on colab