I’m trying to replicate the MNIST example from here
But the code runs for 1 epoch and then gets stuck when it starts the 2nd epoch. It doesn’t throw an error but also doesn’t continue executing.
Please find the colab notebook here
I really need to be able to use TPU for my training, please let me know what I’m doing wrong.
Thanks in Advance